nifty.re.check_model module#

check_model(model, pos)[source]#

Benchmarks and analyzes a NIFTy model’s forward pass, JVP, and VJP.

Runs four types of analysis for each of the three evaluation modes (forward, JVP, VJP):

  • Timing: Measures execution time with and without JIT compilation.

  • Memory analysis: Reports memory usage of the JIT-compiled computations.

  • HLO parsing: Inspects the compiled HLO representation to analyze inlined constants.

  • Model leaves: Inspects the leaves of the model which are not inlined but treated as variables.

Parameters:

  • model (Model or callable) – The NIFTy Model to be benchmarked and analyzed.

  • pos (Any) – Input to the model. Must be a JAX-compatible pytree of arrays.