TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.Runtime.Autograd.Engine.LibTorch.Shape

axisMap generation #

source
def Runtime.Autograd.LibTorch.Broadcast.axisMap {s₁ s₂ : Spec.Shape} (_cb : s₁.CanBroadcastTo s₂) :
Array ℕ

Generate the CUDA axisMap for right-aligned broadcasting.

The proof establishes compatibility, but the map itself is determined solely by the two ranks: missing leading axes map to zero, and every remaining output axis maps to its aligned input axis.

Instances For
    source
    def Runtime.Autograd.LibTorch.Broadcast.afterSumAxisMap (s : Spec.Shape) (axis : ℕ) :
    Array ℕ

    CUDA axis map that restores the axis removed by shapeAfterSum.

    An out-of-range axis leaves the map unchanged, matching shapeAfterSum. Constructing each entry once avoids repeatedly copying and shifting the suffix while descending through the shape.

    Instances For
      source
      def Runtime.Autograd.LibTorch.Broadcast.afterSumArgs (s : Spec.Shape) (axis : ℕ) :
      Array ℕ × Array ℕ × Array ℕ

      CUDA metadata for TorchLean.Tensor.broadcastAfterSum.

      Instances For