TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.Runtime.Autograd.Engine.Cuda.Shape

axisMap generation #

source
def Runtime.Autograd.Cuda.Broadcast.shiftInputAxes (m : Array ℕ) :
Array ℕ

Shift every nonzero entry of a partial axis map up by one, making room for a leading axis.

Instances For
    source
    def Runtime.Autograd.Cuda.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.Cuda.Broadcast.afterSumAxisMap (s : Spec.Shape) (axis : ℕ) :
      Array ℕ

      CUDA axis map that restores the axis removed by shapeAfterSum.

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

        CUDA metadata for TorchLean.Tensor.broadcastAfterSum.

        Instances For