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.

@[inline]

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

Pointwise addition node for two tensors with the same shape.

@[inline]

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

Pointwise subtraction node for two tensors with the same shape.

@[inline]

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

Pointwise multiplication node for two tensors with the same shape.

@[inline]

def Runtime.Autograd.LibTorch.Tape.neg {s : Spec.Shape} (t : Tape SomeBuffer) (xId : ℕ) : Result (Tape SomeBuffer × ℕ)

Direct pointwise negation and its sign-reversing pullback.

@[inline]

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

Multiply by a scalar constant.

@[inline]

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

Pointwise absolute-value node.

@[inline]

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

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

@[inline]

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

Clamp each element to [lo, hi].

@[inline]

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

Pointwise maximum node. Distinct parents split ties according to Buffer.maxBwd. Repeated references contribute the upstream value once, before graph accumulation.

@[inline]

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

Pointwise minimum node. Distinct parents split ties according to Buffer.minBwd. Repeated references use the coalesced-contribution convention of Tape.max.

@[inline]

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

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

@[inline]

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

Pointwise ReLU node with zero derivative on the nonpositive branch.

@[inline]

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

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

@[inline]

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

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.

@[inline]

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

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

@[inline]

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

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

@[inline]

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

Elementwise reciprocal 1/x.

@[inline]

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

Elementwise log(softplus(x) + ε).

The caller must choose ε positive after conversion to binary32 to keep a rounded-zero softplus away from log(0). This operation does not validate ε or guarantee finite results for nonfinite inputs.

@[inline]

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

Direct sigmoid values and the input-based weighted pullback.

@[inline]

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

Direct tanh values and the input-based weighted pullback.

@[inline]

def Runtime.Autograd.LibTorch.Tape.silu {s : Spec.Shape} (t : Tape SomeBuffer) (xId : ℕ) : Result (Tape SomeBuffer × ℕ)

SiLU as one CUDA tape node. ATen evaluates the split-exponential value and local VJP.

@[inline]

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

Tanh-approximate GELU as one CUDA tape node.

The adapter evaluates the staged pointwise formula with ATen operations. TorchLean records the node and owns its VJP rule through Activation.geluDerivSpec; one tape node does not imply one GPU kernel launch.

@[inline]

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

Pointwise softplus node with sigmoid derivative.