jax.core.ClosedJaxpr

class jax.core.ClosedJaxpr(jaxpr, consts)

Parameters:

• jaxpr (Jaxpr)

• consts (Sequence)

__init__(jaxpr, consts)

Parameters:

• jaxpr (Jaxpr)

• consts (Sequence)

Methods

__init__(jaxpr, consts)
map_jaxpr(f)
pretty_print(*[, source_info, print_shapes, ...])
replace(*[, jaxpr, consts])

Attributes

consts
effects
eqns
in_avals
jaxpr
literals
out_avals