TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.Examples.Interop.JAX

From Jaxpr to a permutation-invariance proof #

The Python companion captures sum(x * x) and several other tensor computations. run loads its artifacts through the shared IR checker and compares the native Lean outputs with the recorded JAX outputs. Theorems below concern the real-valued IR computation, for every vector length and every permutation, not just those samples.

Neither numerical agreement nor these theorems certify the Python adapter or XLA. Permutation invariance uses exact addition and does not assert bitwise invariance of a floating-point reduction.

source
def NN.Examples.Interop.JAX.sumSquaresGraph (shape : Spec.Shape) :
IR.Graph

The three IR nodes emitted for jax.numpy.sum(x * x).

Instances For
    source
    noncomputable def NN.Examples.Interop.JAX.sumSquares {shape : Spec.Shape} (x : TorchLean.Tensor ℝ shape) :
    ℝ

    The tensor calculation represented by the captured graph.

    Instances For
      source
      theorem NN.Examples.Interop.JAX.sumSquaresGraph_denote {shape : Spec.Shape} (x : TorchLean.Tensor ℝ shape) :
      (sumSquaresGraph shape).denote { } (Spec.SomeTensor.ofTensor x) 2 = Except.ok (Spec.SomeTensor.ofTensor (TorchLean.Tensor.scalar (sumSquares x)))

      Evaluate the complete IR graph, rather than proving only a separate tensor formula.

      source
      theorem NN.Examples.Interop.JAX.sumSquares_permutation {n : ℕ} (permutation : Equiv.Perm (Fin n)) (x y : TorchLean.Tensor ℝ [n]) (h : ∀ (i : Fin n), y[i] = x[permutation i]) :
      sumSquares y = sumSquares x

      Reordering vector entries leaves the exact real sum of their squares unchanged.

      source
      theorem NN.Examples.Interop.JAX.sumSquaresGraph_permutation {n : ℕ} (permutation : Equiv.Perm (Fin n)) (x y : TorchLean.Tensor ℝ [n]) (h : ∀ (i : Fin n), y[i] = x[permutation i]) :
      (sumSquaresGraph [n]).denote { } (Spec.SomeTensor.ofTensor y) 2 = (sumSquaresGraph [n]).denote { } (Spec.SomeTensor.ofTensor x) 2

      The captured graph itself is permutation invariant under its real-valued semantics.

      source
      def NN.Examples.Interop.JAX.checkBinary (artifact : Lean.Json) :
      IO Unit

      Compare a mixed-precision example using exact binary words at every declared precision.

      This runner checks the recorded example, not all inputs or XLA's reduction schedule. The companion chooses exactly representable reduction inputs for these comparisons.

      Instances For
        source
        def NN.Examples.Interop.JAX.check (artifact : Lean.Json) :
        IO Unit

        Check an artifact and compare all designated outputs using the binary64 diagnostic evaluator.

        Instances For
          source
          def NN.Examples.Interop.JAX.run (paths : List String) :
          IO Unit

          Read one or more artifacts produced by the Python companion.

          Instances For