Checking captured tensor computations

This entry point checks a typed IR artifact and optionally executes supplied arguments or compares supplied reference outputs. The result states the scope of each check. Passing these checks is not a theorem about all inputs, a proof of the frontend translation, or a validation of external kernels.

The report preserves the artifact's capture_constraints as unchecked metadata. Structure, type and example-execution checks do not establish these conditions for all inputs.

compare treats two NaN results as equal whatever their sign and payload, which the IR leaves unspecified; compare_bitwise requires equal words for NaNs too. Any other difference between the IR's results and the reference fails. For a Jaxpr artifact whose reference XLA's CPU backend computed, the checker then runs the program a second time as XLA's compiled CPU code computes it, where this is exactly emulable: reciprocals of broadcast and constant divisors, jax.jit's division and power rewrites when the artifact records a jax.jit reference, and the flushing of subnormal values. A value that XLA's own rsqrt or power function computes counts as reproduced within the window of Internal.window, measured in the format the function computes in. If the second run reproduces every differing element, the failed finding lists the behaviours in consistent_with. That field is advice, not an explanation: an importer error with the same effect on these arguments cannot be ruled out, and some importer errors coincide with XLA's own rewrites. Under jax.jit it also requires the artifact's primitive-by-primitive reference, when there is one, to agree with the IR or to be reproduced without jax.jit's rewrites. The checker trusts the artifact's record of how the reference was computed, so a forged record can at most add this advice to a failure.

Unknown requests are reported as unsupported. They must not disappear from a successful report. Existing certificate checkers and mathematical proof APIs remain separate tools with their own contracts.

def NN.Verification.IR.Internal.finding (name status message : String) (extra : List (String × Lean.Json) := []) : Lean.Json

One named check with its outcome and an explanation of its scope.

def NN.Verification.IR.Internal.value (json : Lean.Json) : Except String IR.Interpreter.reference.Value

Read a supported tensor from complete element words, with no host-float conversion.

def NN.Verification.IR.Internal.binaryWordClass (format : FloatLib.Floats.Formats.BinaryInterchange.FloatFormat) (word : ℕ) : Bool × Bool

FloatLib's classification of one binary word: whether it is a NaN and whether it is subnormal, under the format's exceptional-value encoding.

def NN.Verification.IR.Internal.agree (bitwise : Bool) (dtype : IR.DType) (a b : ℕ) : Bool

Whether two words of dtype agree: equal, or both NaNs unless bitwise is set.

structure NN.Verification.IR.Internal.Mismatch : Type

The elements of one result that disagree with its reference.

• elements : Array ℕ

  Positions of the elements that disagree under the requested NaN policy.

• nanWords : ℕ

  Elements where both words are NaNs with a different sign or payload.

• total : ℕ

  Elements of the result.

def NN.Verification.IR.Internal.compareResult (bitwise : Bool) (output expected : IR.Interpreter.reference.Value) : Except String Mismatch

Compare one result with its reference. Two NaNs agree unless bitwise is set.

XLA's documented CPU behaviour

The behaviours, how the artifact says XLA compiled the reference, and the emulated program that compare runs for a differing Jaxpr result.

structure NN.Verification.IR.Internal.XlaBehaviour : Type

A documented way in which XLA's compiled CPU code computes differently from the IR, with the code that a failed comparison lists in consistent_with and the description of NN/IR/README.md.

• code : String

  Stable code in the report.

• description : String

  What XLA does differently.

def NN.Verification.IR.Internal.instBEqXlaBehaviour.beq : XlaBehaviour → XlaBehaviour → Bool

@[instance_reducible]

instance NN.Verification.IR.Internal.instBEqXlaBehaviour : BEq XlaBehaviour

def NN.Verification.IR.Internal.divisorReciprocal : XlaBehaviour

A division by a broadcast divisor, or under jax.jit by a constant.

def NN.Verification.IR.Internal.rewrittenDivision : XlaBehaviour

The simplifier's division rewrites under jax.jit.

def NN.Verification.IR.Internal.powerProducts : XlaBehaviour

Constant exponents computed as products under jax.jit.

def NN.Verification.IR.Internal.subnormal : XlaBehaviour

Subnormal values flushed to zero.

def NN.Verification.IR.Internal.reciprocalSquareRoot : XlaBehaviour

XLA's own reciprocal square root.

def NN.Verification.IR.Internal.power : XlaBehaviour

XLA's own power function.

def NN.Verification.IR.Internal.window (behaviour : XlaBehaviour) (dtype : IR.DType) : Option ℕ

The window, in units in the last place of dtype, within which a value that XLA's own rsqrt or power function computes in dtype counts as reproduced: the largest distance measured between XLA's function and the IR's in that format, 2 for rsqrt and 1 for the power, in binary32 and binary64. XLA's power is within 1 ulp of the IR's for integral exponents and where |y log x| < 1; elsewhere FloatLib's power is farther from it, and such a result is not reproduced. XLA computes binary16 and bfloat16 rsqrt and powers in binary32 and rounds once, so the window applies to the binary32 value, before that rounding; a node of a reduced format has none.

structure NN.Verification.IR.Internal.Compilation : Type

How XLA compiled a Jaxpr reference, as the artifact records it.

• jit : Bool

  The reference came from jax.jit, which compiles the whole program as one computation.

• computation : Array (Option ℕ)

  Otherwise, for each node, the computation XLA compiled it in: the JAX importer's index of the primitive, or of the nested jax.jit call, that produced it.

• jitComputations : Array ℕ

  The computations that are nested jax.jit calls.

def NN.Verification.IR.Internal.Compilation.read (artifact : Lean.Json) (perPrimitive : Bool := false) : Except String Compilation

Read reference_execution, jit_computations and each node's computation; or, with perPrimitive, describe the primitive-by-primitive reference, where each primitive, inside nested jax.jit calls too, is a computation of its own, identified by the node's primitive_index. A raw node that the lowering expands owns every node up to its result, as in the type assignment of Import.IR.Internal.parse.

def NN.Verification.IR.Internal.Compilation.jitted (compilation : Compilation) (id : ℕ) : Bool

Whether node id is compiled by jax.jit.

def NN.Verification.IR.Internal.Compilation.usesJit (compilation : Compilation) : Bool

Whether jax.jit compiles any part of the program.

structure NN.Verification.IR.Internal.Emulation : Type

The program as XLA's CPU code computes it, built node by node from the IR program.

• jit : Bool

  The reference came from jax.jit.

• nodes : Array IR.Node

  Nodes of the emulated program.

• dtypes : Array IR.DType

  Element type of each node.

• primitives : Array (Option IR.Primitive)

  Casts and registered operations, aligned with nodes.

• constants : Array (Option IR.Interpreter.reference.Value)

  Values of the copied constants and of the ones that the rewrites add.

• origin : Array (Option ℕ)

  The IR node each copied node comes from, for its payload entries.

• computation : Array (Option ℕ)

  The computation each node belongs to, as in Compilation.

• folded : Array Bool

  Whether jax.jit folds each node at compile time: it is computed from constants alone within one jax.jit computation.

• site : Array (Option (XlaBehaviour × ℕ))

  The documented behaviour that produced each node, with the IR node it rewrites.

def NN.Verification.IR.Internal.Emulation.within (emulation : Emulation) (computation : Option ℕ) (id : ℕ) : Bool

Whether node id belongs to computation. Under jax.jit every node does.

def NN.Verification.IR.Internal.Emulation.isFolded (emulation : Emulation) (id : ℕ) : Bool

Whether node id is folded at compile time.

def NN.Verification.IR.Internal.Emulation.operation? (emulation : Emulation) (computation : Option ℕ) (id : ℕ) : Option (IR.OpKind × Array ℕ)

The operation and operands of node id when it belongs to computation.

def NN.Verification.IR.Internal.Emulation.broadcastIn (emulation : Emulation) (computation : Option ℕ) (id : ℕ) : Bool

Whether a divisor is a broadcast computed in computation, past reshapes and permutations.

structure NN.Verification.IR.Internal.Site : Type

The IR node being emulated, its computation, and the documented behaviour that the new nodes emulate, if any.

• node : ℕ

  The IR node.

• computation : Option ℕ

  Its computation.

• jitted : Bool

  Whether jax.jit compiles it.

• difference : Option XlaBehaviour

  The behaviour the new nodes emulate.

@[reducible, inline]

abbrev NN.Verification.IR.Internal.Build (α : Type) : Type

Construction of the emulated program.

def NN.Verification.IR.Internal.emit (kind : IR.OpKind) (parents : Array ℕ) (shape : Spec.Shape) (dtype : IR.DType) (site : Site) (primitive : Option IR.Primitive := none) (constant : Option IR.Interpreter.reference.Value := none) (copy : Bool := false) : Build ℕ

Add a node for site, returning its id.

def NN.Verification.IR.Internal.integerWord (format : FloatLib.Floats.Formats.BinaryInterchange.FloatFormat) (n : ℤ) : ℕ

The word of the integer n rounded once to format.

def NN.Verification.IR.Internal.ones (dtype : IR.DType) (shape : Spec.Shape) (site : Site) : Build ℕ

A constant of ones.

def NN.Verification.IR.Internal.castTo (id : ℕ) (target : IR.DType) (site : Site) : Build ℕ

Round node id to target.

def NN.Verification.IR.Internal.rewritesIn (dtype : IR.DType) : Bool

Formats in which jax.jit rewrites quotients of quotients and quotients by quotients, each step rounded to the format. XLA computes bfloat16 in binary32, where the conversions between operations block the rewrites.

def NN.Verification.IR.Internal.rootRewritesIn (dtype : IR.DType) : Bool

Formats in which the emulation rewrites divisions by square roots, reciprocal square roots and powers. The IR computes binary16 rsqrt and pow in binary32, so these patterns are not binary16 operations there.

def NN.Verification.IR.Internal.emitDivision (fuel num den : ℕ) (shape : Spec.Shape) (dtype : IR.DType) (site : Site) : Build ℕ

A division as XLA's CPU code computes it. A broadcast divisor, or under jax.jit a constant one, becomes multiplication by its reciprocal rounded to the format; bfloat16 computes both in binary32 and rounds the product once. Otherwise, under jax.jit, the simplifier's rewrites (x / y) / z = x / (y * z) and x / (y / z) = (x * z) / y apply in rewritesIn formats, then x / sqrt y = x * rsqrt y, x / rsqrt y = x * sqrt y and x / y ^ c = x * y ^ -c in rootRewritesIn formats, in that order, and the division a rewrite produces is computed by the same rules.

def NN.Verification.IR.Internal.productExponent (dtype : IR.DType) (exponent : IR.Interpreter.reference.Value) : Option ℤ

The exponent 2, 3 or -1 that jax.jit computes as products or a reciprocal, when every element of a binary16, binary32 or binary64 exponent is that integer.

def NN.Verification.IR.Internal.halfSource (emulation : Emulation) (id : ℕ) : Option ℕ

The binary16 operand of a binary32 node: a conversion of a binary16 value, as the JAX importer writes the base of a binary16 power.

def NN.Verification.IR.Internal.emulate (program : IR.Program) (compilation : Compilation) (values : Array IR.Interpreter.reference.Value) : Except String (Emulation × Array ℕ)

Build the emulated program from the IR program and its node values, returning the emulated node of each IR node. Under jax.jit a constant exponent 2, 3 or -1 becomes products or a reciprocal, each rounded to the format; XLA rounds every product of a binary16 base to binary16, so the binary32 power that the JAX importer writes for it is computed in binary16 and widened.

def NN.Verification.IR.Internal.flushes (dtype : IR.DType) : Bool

Whether XLA's CPU code flushes subnormal values of dtype: bfloat16, binary32 and binary64. It computes binary16 in binary32, where binary16 subnormals are normal.

def NN.Verification.IR.Internal.flushingOperation (kind : IR.OpKind) (dtype : IR.DType) : Bool

Operations whose subnormal operands and results XLA's CPU code flushes to zero, as measured. Negation, abs, sin, layout operations and binary32 and bfloat16 tanh keep them.

def NN.Verification.IR.Internal.flushing (emulation : Emulation) (node : IR.Node) : Bool × Bool

Whether a node flushes its operands and its result. A node that jax.jit folds flushes neither; a conversion from binary32 or bfloat16 to binary64 flushes its operand, and one from binary64 to binary32 or bfloat16 its result.

def NN.Verification.IR.Internal.flushValue (value : IR.Interpreter.reference.Value) : Except String (IR.Interpreter.reference.Value × Bool)

Replace the subnormal elements of a bfloat16, binary32 or binary64 value by zero of the same sign, reporting whether any element changed.

def NN.Verification.IR.Internal.emulationPayload (emulation : Emulation) (payload : IR.Interpreter.reference.Payload) : IR.Interpreter.reference.Payload

The IR payload, read through the origins of the emulated nodes.

structure NN.Verification.IR.Internal.Emulated : Type

The emulated program and its interface, as a Program, with the payload read through it.

• emulation : Emulation

  The documented behaviours applied to each node.

• program : IR.Program

  The emulated program.

• payload : IR.Interpreter.reference.Payload

  Its payload.

• arguments : Array IR.Interpreter.reference.Value

  The arguments.

def NN.Verification.IR.Internal.Emulated.evalNode (emulated : Emulated) (values : Array IR.Interpreter.reference.Value) (node : IR.Node) : Except String (IR.Interpreter.reference.Value × Bool)

Evaluate one emulated node on values, flushing its operands and its result as flushing says, and report whether flushing changed a word.

def NN.Verification.IR.Internal.Emulated.run (emulated : Emulated) : Except String (Array IR.Interpreter.reference.Value × Array Bool)

Run the emulated program, returning every node's value and whether flushing changed a word of it.

structure NN.Verification.IR.Internal.Approximation : Type

A result that XLA's own function computes, and the window within which it counts as reproduced.

• behaviour : XlaBehaviour

  XLA's function.

• source : ℕ

  The emulated rsqrt or pow node.

• ulps : Except String ℕ

  The window, in units in the last place of the source node's format, or why there is none.

• path : Array ℕ

  The nodes from the source to the result, in order, the result last.

def NN.Verification.IR.Internal.exactArgument (emulation : Emulation) (id : ℕ) : Bool

Whether node id is an argument, reached through layout operations and conversions that do not round.

def NN.Verification.IR.Internal.approximation (emulation : Emulation) (id : ℕ) : Option Approximation

XLA's own function computing result id: an rsqrt or pow node, reached through layout operations, negation and conversions between binary formats, with its window when the node is of binary32 or binary64. A power has a window only when both its operands are arguments (exactArgument), so that XLA's function and the IR's receive the same values: a power with a small exponent hides a large error in its base, and one with a large base a small error in its exponent.

def NN.Verification.IR.Internal.signBit (format : FloatLib.Floats.Formats.BinaryInterchange.FloatFormat) (word : ℕ) : Bool

Whether a word of format has its sign bit set.

def NN.Verification.IR.Internal.isInfinite (format : FloatLib.Floats.Formats.BinaryInterchange.FloatFormat) (word : ℕ) : Bool

Whether a word of format is an infinity.

def NN.Verification.IR.Internal.windowEnds (format : FloatLib.Floats.Formats.BinaryInterchange.FloatFormat) (ulps word : ℕ) : ℕ × ℕ

The words ulps steps below and above word in the order of the magnitudes of an IEEE format, keeping its sign and staying finite: the window never crosses zero or reaches an infinity. An infinity or a NaN is its own window.

def NN.Verification.IR.Internal.windowAt (emulated : Emulated) (values : Array IR.Interpreter.reference.Value) (approximation : Approximation) : Except String (Option (Array ℕ × Array ℕ))

The ends of the window around each element of the approximation's source node, carried to the result through its path: rounding is monotone, and negation and layout operations move the ends with the value.

def NN.Verification.IR.Internal.inWindow (format : FloatLib.Floats.Formats.BinaryInterchange.FloatFormat) (emulated wanted low high : ℕ) : Bool

Whether the reference word wanted lies between the window's ends low and high at the result and has the sign and the finiteness of the emulated word: a difference in the sign of zero or between a finite value and an infinity is outside every window.

def NN.Verification.IR.Internal.dependencies (nodes : Array IR.Node) (output : ℕ) : Array Bool

The nodes that node output depends on, itself included.

def NN.Verification.IR.Internal.hex (word : ℕ) : String

A word in hexadecimal.

structure NN.Verification.IR.Internal.WindowUse : Type

The window used for a result: XLA's function, its width and format, and the differing elements within it.

• behaviour : XlaBehaviour

  XLA's function.

• ulps : ℕ

  Units in the last place.

• dtype : IR.DType

  The format the function computes in.

• result : IR.DType

  The result's format, which the window's ends are rounded to.

• elements : ℕ

  Differing elements within the window.

structure NN.Verification.IR.Internal.Reproduction : Type

What the emulation of XLA's CPU code reproduces of one differing result.

• result : ℕ

  Position of the result.

• mismatch : Mismatch

  Its differing elements.

• applied : Array (XlaBehaviour × ℕ)

  Documented behaviours applied to nodes the result depends on, each with the IR nodes it rewrites, or for flushing the emulated nodes where it changed a word.

• reproduced : ℕ

  Differing elements whose words the emulated program returns.

• window : Option WindowUse

  The window of XLA's own function computing the result, if one applies.

• withheld : Option (XlaBehaviour × String)

  XLA's own function computing the result when no window applies, and why.

• notReproduced : Array (ℕ × ℕ × ℕ × ℕ)

  Elements not reproduced, with the IR's, the emulation's and the reference's words.

def NN.Verification.IR.Internal.Reproduction.used (entry : Reproduction) : Array XlaBehaviour

The behaviours the reproduction used: those applied, and XLA's function if its window holds a differing element.

def NN.Verification.IR.Internal.Reproduction.toJson (entry : Reproduction) : Lean.Json

The report entry of a differing result.

def NN.Verification.IR.Internal.Reproduction.behaviours (entry : Reproduction) : String

The behaviours applied, in words.

def NN.Verification.IR.Internal.Reproduction.windowText (use : WindowUse) : String

The window, in words.

def NN.Verification.IR.Internal.Reproduction.describe (entry : Reproduction) : String

One sentence for the report message.

def NN.Verification.IR.Internal.reproduceResult (bitwise : Bool) (emulated : Emulated) (changed : Array Bool) (values : Array IR.Interpreter.reference.Value) (output index : ℕ) (ir reference : IR.Interpreter.reference.Value) (mismatch : Mismatch) : Except String Reproduction

What the emulated program reproduces of result index's differing elements.

def NN.Verification.IR.Internal.reproduce (bitwise : Bool) (compilation : Compilation) (program : IR.Program) (payload : IR.Interpreter.reference.Payload) (arguments values results reference : Array IR.Interpreter.reference.Value) (differing : Array (ℕ × Mismatch)) : Except String (Array Reproduction)

Run the emulation of XLA's CPU code that compilation describes and report what it reproduces of each differing result.

def NN.Verification.IR.Internal.mismatches (bitwise : Bool) (artifact : Lean.Json) (field : String) (results : Array IR.Interpreter.reference.Value) : Except String (Array IR.Interpreter.reference.Value × Array (ℕ × Mismatch))

Compare the IR's results with the reference tensors in field, returning the references and the elements that differ.

def NN.Verification.IR.Internal.usedBehaviours (entries : Array Reproduction) : Array XlaBehaviour

The behaviours a set of reproductions used, each once, in order.

def NN.Verification.IR.Internal.behaviourCodes (entries : Array Reproduction) : Array String

Their codes.

def NN.Verification.IR.Internal.eagerCheck (bitwise : Bool) (artifact : Lean.Json) (program : IR.Program) (payload : IR.Interpreter.reference.Payload) (arguments values results : Array IR.Interpreter.reference.Value) : Bool × String × Lean.Json

What the primitive-by-primitive reference says about a mismatch with a jax.jit reference: whether the IR agrees with it, or the emulation of XLA's CPU code without jax.jit reproduces it, with a sentence and report entries.

def NN.Verification.IR.Internal.compareFinding (name : String) (artifact : Lean.Json) (program : IR.Program) (payload : IR.Interpreter.reference.Payload) (arguments values results : Array IR.Interpreter.reference.Value) : Bool × Lean.Json

The compare or compare_bitwise finding for computed results and every node value. Any difference fails. For a Jaxpr artifact with a CPU reference, a failure whose differing elements the emulation of XLA's CPU code reproduces carries the behaviours in consistent_with, as advice: an importer error with the same effect cannot be ruled out. With a jax.jit reference, the artifact's primitive-by-primitive reference, when it has one, must also agree with the IR or be reproduced without jax.jit.

def NN.Verification.IR.Internal.result (value : IR.Interpreter.reference.Value) : Except String Lean.Json

Serialize computed results for inspection without a lossy numerical display conversion.

def NN.Verification.IR.check (artifact : Lean.Json) (requests : Array String := #["structure", "types"]) : Lean.Json

Run requested checks, returning explicit failures and unsupported requests as report entries.

execute, compare and compare_bitwise concern only arguments stored in the artifact. A comparison uses complete binary words, except that compare lets two NaNs agree whatever their sign and payload. It can expose legitimate differences in reduction or elementary-function implementations; a mismatch alone does not identify a bug, and it always fails. For a Jaxpr artifact, the failure lists in consistent_with the documented behaviours of XLA's CPU code whose emulation reproduces every differing element (see the module documentation).

def NN.Verification.IR.run (arguments : List String) : IO Unit

CLI: check an artifact as JSON, or emit the existing PyTorch capture adapter.