liveBook · Manning

Jax.make_jaxpr()

As explained in Deep Learning with JAX

A JAX transformation that returns the Jaxpr corresponding to a Python function call signature.

Need to see what JAX will transform before sending computation to XLA? Call jax.make_jaxpr() on a function and invoke the returned transformation with representative arguments. It exposes Jaxpr, a typed intermediate representation containing inputs, constants, primitive equations, outputs, shapes, and dtypes. section 5.2.1

Inspecting a reduction

For f1(x, y, z) = jnp.sum(x + y * z), jax.make_jaxpr(f1)(x, y, z) reveals broadcasting, multiplication, addition, and reduce_sum operations. The printed variables expose array shapes such as f32[3,3] and the scalar result f32[].

How Deep Learning with JAX covers it

How deep it goes
Displays typed variables, constants, primitive equations, shapes, and control-flow sub-Jaxprs, but does not execute compiled machine code.
Examples
Demonstrates reductions, side effects, fixed and dynamic loops, dynamic conditionals, static arguments, lax.fori_loop, and lax.cond.
What you can do after
Diagnose tracing behavior, inspect staged operations, and identify control-flow errors before using JIT compilation.
How it connects
Connects Jaxpr with tracing, JIT compilation, XLA, primitive operations, and structured control flow.
Cover of Deep Learning with JAX
liveBookmanning.com
Read Deep Learning with JAX on liveBook Get the book at manning.com

Related concepts

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.