ClosedJaxpr(jaxpr, consts)

Jaxpr(constvars, invars, outvars, eqns[, ...])

JaxprEqn(invars, outvars, primitive, params, ...)

Literal(val, aval)

Primitive(name)

Token(buf)

Var(aval[, initial_qdd, final_qdd])

array_types

set() -> new empty set object set(iterable) -> new set object

jaxpr_as_fun

primitives

mapped_aval(size, axis, aval)

unmapped_aval(size, axis, aval[, ...])