TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.Runtime.Autograd.IRExec.Lowering.Reductions

Reduction IR Lowering #

Checked lowering for broadcasts, axis reductions, full reductions, and scalar losses.

Each operation has its own small lower* definition. lowerReduction only dispatches on the operation kind, and the lowerReduction_* equation lemmas let correctness proofs reduce a dispatch to the branch they care about without unfolding the whole dispatcher.

source
def Runtime.Autograd.IRExec.Internal.lowerBroadcastTo {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (s₁ s₂ : Spec.Shape) :
NodeLoweringResult ctx

Checked lowering for .broadcastTo s₁ s₂.

Instances For
    source
    def Runtime.Autograd.IRExec.Internal.lowerReduceSum {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (axis : ℕ) :
    NodeLoweringResult ctx

    Checked lowering for .reduceSum axis.

    Instances For
      source
      def Runtime.Autograd.IRExec.Internal.lowerReduceMean {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (axis : ℕ) :
      NodeLoweringResult ctx

      Checked lowering for .reduceMean axis.

      Instances For
        source
        def Runtime.Autograd.IRExec.Internal.lowerSum {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) :
        NodeLoweringResult ctx

        Checked lowering for .sum.

        Instances For
          source
          def Runtime.Autograd.IRExec.Internal.lowerMseLoss {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) :
          NodeLoweringResult ctx

          Checked lowering for .mseLoss.

          Instances For
            source
            def Runtime.Autograd.IRExec.Internal.lowerReduction {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (kind : NN.IR.OpKind) :
            NodeLoweringResult ctx

            Checked lowering for broadcasts, axis reductions, full reductions, and scalar losses.

            Instances For
              source
              @[simp]
              theorem Runtime.Autograd.IRExec.Internal.lowerReduction_broadcastTo {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (s₁ s₂ : Spec.Shape) :
              lowerReduction ctx (NN.IR.OpKind.broadcastTo s₁ s₂) = lowerBroadcastTo ctx s₁ s₂

              Dispatch equation for .broadcastTo s₁ s₂.

              source
              @[simp]
              theorem Runtime.Autograd.IRExec.Internal.lowerReduction_reduceSum {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (axis : ℕ) :
              lowerReduction ctx (NN.IR.OpKind.reduceSum axis) = lowerReduceSum ctx axis

              Dispatch equation for .reduceSum axis.

              source
              @[simp]
              theorem Runtime.Autograd.IRExec.Internal.lowerReduction_reduceMean {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) (axis : ℕ) :
              lowerReduction ctx (NN.IR.OpKind.reduceMean axis) = lowerReduceMean ctx axis

              Dispatch equation for .reduceMean axis.

              source
              @[simp]
              theorem Runtime.Autograd.IRExec.Internal.lowerReduction_sum {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) :
              lowerReduction ctx NN.IR.OpKind.sum = lowerSum ctx

              Dispatch equation for .sum.

              source
              @[simp]
              theorem Runtime.Autograd.IRExec.Internal.lowerReduction_mseLoss {α : Type} [TorchLean.Storage α] [Context α] {Γ : List Spec.Shape} (ctx : NodeLoweringContext α Γ) :
              lowerReduction ctx NN.IR.OpKind.mseLoss = lowerMseLoss ctx

              Dispatch equation for .mseLoss.