jax.extend.core.Jaxpr#

class jax.extend.core.Jaxpr(constvars, invars, outvars, eqns, effects=frozenset({}), debug_info=None, is_high=False)[source]#
Parameters:

  • constvars (Sequence[Var])

  • invars (Sequence[Var])

  • outvars (Sequence[Atom])

  • eqns (Sequence[JaxprEqn])

  • effects (Effects)

  • debug_info (DebugInfo)

  • is_high (bool)

__init__(constvars, invars, outvars, eqns, effects=frozenset({}), debug_info=None, is_high=False)[source]#
Parameters:

  • constvars (Sequence[Var]) – list of variables introduced for constants. Array constants are replaced with such variables while scalar constants are kept inline.

  • invars (Sequence[Var]) – list of input variables. Together, constvars and invars are the inputs to the Jaxpr.

  • outvars (Sequence[Atom]) – list of output atoms.

  • eqns (Sequence[JaxprEqn]) – list of equations.

  • effects (Effects) – set of effects. The effects on a jaxpr are a superset of the union of the effects for each equation.

  • debug_info (DebugInfo) – debugging information.

  • is_high (bool)

Methods

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

pretty_print(*[, source_info, print_shapes, ...])

replace(**kwargs)

Attributes

constvars

debug_info

effects

eqns

invars

is_high

outvars