CUDA Tape Operations: Elementwise Nodes

Elementwise ops

The backward closures below return newly allocated gradient buffers. When a derivative uses intermediate CUDA buffers, it releases those intermediates before returning the final gradient. The returned buffers are owned by the tape/gradient accumulator; workspace buffers are owned locally.

def Runtime.Autograd.Cuda.Tape.add {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Pointwise addition node for two tensors with the same shape.

def Runtime.Autograd.Cuda.Tape.sub {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Pointwise subtraction node for two tensors with the same shape.

def Runtime.Autograd.Cuda.Tape.mul {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Pointwise multiplication node for two tensors with the same shape.

def Runtime.Autograd.Cuda.Tape.scale {s : Spec.Shape} (t : Tape) (xId : ℕ) (c : Float) : Result (Tape × ℕ)

Multiply by a scalar constant.

def Runtime.Autograd.Cuda.Tape.abs {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Pointwise absolute-value node.

def Runtime.Autograd.Cuda.Tape.sqrt {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Pointwise square-root node using the CUDA buffer derivative convention.

def Runtime.Autograd.Cuda.Tape.clamp {s : Spec.Shape} (t : Tape) (xId : ℕ) (lo hi : Float) : Result (Tape × ℕ)

Clamp each element to [lo, hi].

def Runtime.Autograd.Cuda.Tape.max {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Pointwise maximum node; the backward rule splits ties according to Buffer.maxBwd.

def Runtime.Autograd.Cuda.Tape.min {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Pointwise minimum node; the backward rule splits ties according to Buffer.minBwd.

def Runtime.Autograd.Cuda.Tape.div {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Pointwise division node with the usual quotient-rule backward closure.

def Runtime.Autograd.Cuda.Tape.relu {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Pointwise ReLU node with zero derivative on the nonpositive branch.

def Runtime.Autograd.Cuda.Tape.exp {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Pointwise exponential node; backward recomputes exp x as local workspace.

def Runtime.Autograd.Cuda.Tape.sin {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise sine with VJP cos(x) * dLdy.

The cosine buffer belongs to this backward call and is released after multiplication. The input and upstream gradient remain owned by the tape.

def Runtime.Autograd.Cuda.Tape.cos {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise cosine with VJP -sin(x) * dLdy, releasing both temporary buffers.

def Runtime.Autograd.Cuda.Tape.log {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Pointwise natural-log node; callers are responsible for the positive-domain convention.

def Runtime.Autograd.Cuda.Tape.inv {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise reciprocal 1/x.

def Runtime.Autograd.Cuda.Tape.safeLog {s : Spec.Shape} (t : Tape) (xId : ℕ) (ε : Float) : Result (Tape × ℕ)

Elementwise "safe log" that protects against log(0) by adding a small ε internally.

Spec semantics: log(softplus(x) + ε).

def Runtime.Autograd.Cuda.Tape.sigmoid {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise sigmoid (logistic).

def Runtime.Autograd.Cuda.Tape.tanh {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Pointwise hyperbolic tangent node.

def Runtime.Autograd.Cuda.Tape.gelu {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Tanh-approximate GELU as one CUDA tape node.

The native kernels fuse only the pointwise numerical work. TorchLean still records the node and owns its VJP rule through Activation.geluDerivSpec.

def Runtime.Autograd.Cuda.Tape.softplus {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Pointwise softplus node with sigmoid derivative.