TorchLean API

NN.Runtime.Autograd.Torch.LinkedSession.ShapeIndex

Proof-Linked Session: Shape and Indexing Operations #

def Runtime.Autograd.Torch.Internal.SessionIR.maxPool {α : Type} (s : SessionIR α) [Context α] [DecidableEq Spec.Shape] {d C : } {inSpatial kernel stride padding : Vector d} {hKernel : ∀ (i : Fin d), kernel.get i 0} (x : TensorRef α (Spec.Shape.ofList (C :: inSpatial.toList))) :
IO (TensorRef α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

N-D max-pooling for channels-first tensors (C, spatial...) (no batch axis).

PyTorch comparison: torch.nn.functional.max_pool1d / max_pool2d / max_pool3d depending on the spatial rank d.

Instances For
    def Runtime.Autograd.Torch.Internal.SessionIR.smoothMaxPool {α : Type} (s : SessionIR α) [Context α] [DecidableEq Spec.Shape] {d C : } {inSpatial kernel stride padding : Vector d} {hKernel : ∀ (i : Fin d), kernel.get i 0} (x : TensorRef α (Spec.Shape.ofList (C :: inSpatial.toList))) (beta : α) :
    IO (TensorRef α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

    N-D smooth max-pooling (log-sum-exp surrogate) for channels-first tensors (C, spatial...).

    This is a differentiable approximation of max-pooling; there is no direct PyTorch primitive.

    Instances For
      def Runtime.Autograd.Torch.Internal.SessionIR.avgPool {α : Type} (s : SessionIR α) [Context α] [DecidableEq Spec.Shape] {d C : } {inSpatial kernel stride padding : Vector d} (hKernel : ∀ (i : Fin d), kernel.get i 0) (x : TensorRef α (Spec.Shape.ofList (C :: inSpatial.toList))) :
      IO (TensorRef α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

      N-D average-pooling for channels-first tensors (C, spatial...) (no batch axis).

      PyTorch comparison: torch.nn.functional.avg_pool1d / avg_pool2d / avg_pool3d depending on the spatial rank d.

      Instances For
        def Runtime.Autograd.Torch.Internal.SessionIR.maxPool2d {α : Type} (s : SessionIR α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : } {h1 : kH 0} {h2 : kW 0} (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) :
        IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

        2D max-pooling for channel-first images.

        PyTorch comparison: torch.nn.functional.max_pool2d (for NCHW-like layouts, here without batch).

        Instances For
          def Runtime.Autograd.Torch.Internal.SessionIR.smoothMaxPool2d {α : Type} (s : SessionIR α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : } {h1 : kH 0} {h2 : kW 0} (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) (beta : α) :
          IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

          Smooth approximation of max-pooling (softmax pooling) for channel-first images.

          This is not a standard PyTorch primitive; conceptually it behaves like applying a softmax over each pooling window with inverse-temperature beta and returning the expected value.

          Instances For
            def Runtime.Autograd.Torch.Internal.SessionIR.avgPool2d {α : Type} (s : SessionIR α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : } (h1 : kH 0) (h2 : kW 0) (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) :
            IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

            2D average-pooling for channel-first images.

            PyTorch comparison: torch.nn.functional.avg_pool2d (for NCHW-like layouts, here without batch).

            Instances For
              def Runtime.Autograd.Torch.Internal.SessionIR.relu {α : Type} (s : SessionIR α) [Mul α] [Add α] [Zero α] [Max α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) :
              IO (TensorRef α sh)

              Record elementwise ReLU.

              PyTorch comparison: torch.relu(x) / torch.nn.functional.relu(x).

              Instances For

                Flatten a tensor into a 1D vector of length Shape.size sh.

                PyTorch comparison: torch.flatten(x) (with default start_dim=0).

                Instances For
                  def Runtime.Autograd.Torch.Internal.SessionIR.reshape {α : Type} (s : SessionIR α) [Inhabited α] [Zero α] [DecidableEq Spec.Shape] {sh1 sh2 : Spec.Shape} (x : TensorRef α sh1) (h : sh1.size = sh2.size) :
                  IO (TensorRef α sh2)

                  Reshape a tensor while preserving total number of elements.

                  The proof argument h enforces Shape.size sh1 = Shape.size sh2. PyTorch comparison: torch.reshape(x, new_shape) / x.view(new_shape) (when contiguous).

                  Instances For

                    Transpose a 2D matrix (swap the two axes).

                    PyTorch comparison: x.t() for 2D tensors, or x.transpose(0, 1).

                    Instances For

                      Permute a 3D tensor by moving the first axis to the end: (a,b,c) → (b,c,a).

                      PyTorch comparison: x.permute(1,2,0) for a 3D tensor.

                      Instances For

                        Permute a 3D tensor by moving the last axis to the front: (a,b,c) → (c,a,b).

                        PyTorch comparison: x.permute(2,0,1) for a 3D tensor.

                        Instances For

                          Swap the last two axes of a 3D tensor: (a,b,c) → (a,c,b).

                          PyTorch comparison: x.transpose(1,2) for a 3D tensor.

                          Instances For

                            Swap two adjacent axes at a given depth inside the shape.

                            This is a more general permutation helper used in some shape-manipulating models. PyTorch comparison: like x.transpose(dim, dim+1) for a suitably chosen dim.

                            Instances For
                              def Runtime.Autograd.Torch.Internal.SessionIR.broadcastTo {α : Type} (s : SessionIR α) [Inhabited α] [Add α] [Zero α] [DecidableEq Spec.Shape] {sh1 sh2 : Spec.Shape} (cb : sh1.CanBroadcastTo sh2) (x : TensorRef α sh1) :
                              IO (TensorRef α sh2)

                              Broadcast a tensor to a larger shape.

                              The witness cb : Shape.CanBroadcastTo sh1 sh2 encodes the broadcasting proof. PyTorch comparison: x.expand(...) / implicit broadcasting.

                              Instances For

                                Sum-reduce along axis.

                                PyTorch comparison: torch.sum(x, dim=axis).

                                Instances For

                                  Mean-reduce along axis.

                                  PyTorch comparison: torch.mean(x, dim=axis).

                                  Instances For

                                    Gather a single scalar x[i] from a 1D vector, with a compile-time Fin n index.

                                    PyTorch comparison: x[i] for a 1D tensor.

                                    Instances For

                                      Gather a row x[i] from a 2D tensor, with a compile-time Fin rows index.

                                      PyTorch comparison: x[i] for a 2D tensor (row indexing).

                                      Instances For

                                        Read a Nat from the nat-environment.

                                        Out-of-bounds reads return 0 (total function), which is convenient for modeling "possibly invalid" indices without throwing.

                                        Instances For

                                          Read a length-k vector of Nats starting at start from the nat-environment.

                                          Out-of-bounds reads fall back to 0 elementwise via natAt.

                                          Instances For

                                            Dynamic gather of a scalar from a 1D vector using a runtime NatRef index.

                                            Out-of-range indices produce 0 instead of raising. PyTorch comparison: similar to x[i] where i is a Python integer, except PyTorch raises on out-of-range while this definition totalizes the behavior for ease of reasoning.

                                            Instances For

                                              Dynamic gather of a row from a 2D tensor using a runtime NatRef index.

                                              Out-of-range indices yield a zero row. PyTorch comparison: similar to x[i] for 2D tensors with runtime i, but PyTorch raises on out-of-range whereas this definition is totalized for ease of reasoning.

                                              Instances For

                                                Dynamic gather of k scalars from a 1D tensor using a runtime NatVecRef k of indices.

                                                Out-of-range indices yield 0. In the VJP, gradients are accumulated for repeated indices (i.e. it behaves like a gather followed by a scatter-add back into the source vector). PyTorch comparison: related to torch.gather / advanced indexing, but with totalized out-of-range behavior.

                                                Instances For

                                                  Dynamic gather of k rows from a 2D tensor using a runtime NatVecRef k of row indices.

                                                  Out-of-range indices yield zero rows. In the VJP, gradients are accumulated into the selected rows (scatter-add semantics), including accumulation for repeated indices. PyTorch comparison: similar to torch.index_select(x, dim=0, index=...) or advanced indexing on the first dimension, but with totalized out-of-range behavior.

                                                  Instances For

                                                    Gather a scalar from a 1D vector using a raw Nat index.

                                                    PyTorch comparison: like x[i] with an integer index, but this operation is recorded into the proved IR (so it is stable for compilation/verification).

                                                    Instances For

                                                      Gather k scalars from a 1D vector using an explicit index tensor.

                                                      PyTorch comparison: related to torch.gather / advanced indexing with an integer index tensor.

                                                      Instances For

                                                        Gather k rows from a 2D tensor using an explicit index tensor.

                                                        PyTorch comparison: similar to torch.index_select(x, dim=0, index=...) or advanced indexing.

                                                        Instances For

                                                          Scatter-add into a vector: return a copy of x with x[i] += v.

                                                          PyTorch comparison: similar to x.scatter_add_(dim=0, index=..., src=...) in spirit, but this is functional (returns a new tensor) and uses a single Fin n index.

                                                          Instances For

                                                            Scatter-add into a matrix row: return a copy of x with x[i, :] += v.

                                                            PyTorch comparison: like adding a row vector into a selected row (functional analogue of an in-place indexed add).

                                                            Instances For