Shape and axis operations

def Runtime.Autograd.Model.F.permute? {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axes : Array ℕ) (x : Ref s) : m (Option ((s' : Spec.Shape) × Ref s'))

Dynamic permutation: like permute, but returns an existential output shape.

PyTorch analogue: torch.permute / Tensor.permute (with runtime checks).

def Runtime.Autograd.Model.F.permute {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s sOut : Spec.Shape} (axes : Array ℕ) (x : Ref s) : m (Option (Ref sOut))

Permutation with an expected output shape.

This calls permute? and checks that the computed shape equals sOut.

def Runtime.Autograd.Model.F.transpose {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s sOut : Spec.Shape} (axis₁ axis₂ : ℕ) (x : Ref s) : m (Option (Ref sOut))

Exchange two arbitrary axes and check the statically expected output shape.

The result is none when either axis is invalid or the transposed shape is not sOut.

def Runtime.Autograd.Model.F.Internal.reduceAlongLast {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] (operation : {s : Spec.Shape} → (axis : ℕ) → [Spec.Shape.HasNonemptyAxis axis s] → [s.WellFormed] → Ref s → m (Ref (TorchLean.Tensor.shapeAfterSum s axis))) (x : (s : Spec.Shape) × Ref s) : m (Option ((s' : Spec.Shape) × Ref s'))

Reduce the last axis after checking that the tensor has nonempty dimensions.

def Runtime.Autograd.Model.F.Internal.reduceAxes {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] (reduceLast : (s : Spec.Shape) × Ref s → m (Option ((s' : Spec.Shape) × Ref s'))) {s : Spec.Shape} (axes : Array ℕ) (keepdim : Bool) (x : Ref s) : m (Option ((s' : Spec.Shape) × Ref s'))

Core implementation for reductions over a runtime array of axes.

This lowers “reduce along axis k” to:

1. permute axis k to the last position,
2. call reduceLast, and
3. optionally re-insert a singleton dimension when keepdim = true.

reduceSumDims? and reduceMeanDims? specialize this operation.

def Runtime.Autograd.Model.F.reduceSumDims? {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axes : Array ℕ) (x : Ref s) (keepdim : Bool := false) : m (Option ((s' : Spec.Shape) × Ref s'))

Dynamic multi-axis sum reduction (like torch.sum(x, dim=axes, keepdim=...)).

def Runtime.Autograd.Model.F.reduceMeanDims? {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axes : Array ℕ) (x : Ref s) (keepdim : Bool := false) : m (Option ((s' : Spec.Shape) × Ref s'))

Dynamic multi-axis mean reduction (like torch.mean(x, dim=axes, keepdim=...)).

def Runtime.Autograd.Model.F.softmax {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (x : Ref s) : m (Ref s)

Softmax along any valid tensor dimension.

The selected dimension is moved to the end for the backend's row-softmax primitive and then moved back. The reverse-swap theorem makes the result shape exactly s; no runtime shape check is needed.

def Runtime.Autograd.Model.F.logSoftmax {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (x : Ref s) : m (Ref s)

Log-softmax along any valid tensor dimension.