jax.core.Jaxpr

class jax.core.Jaxpr(constvars, invars, outvars, eqns, effects={})

Parameters

• constvars (Sequence[Var]) –

• invars (Sequence[Var]) –

• outvars (Sequence[Union[Var, Literal]]) –

• eqns (Sequence[JaxprEqn]) –

• effects (Set[Hashable]) –

__init__(constvars, invars, outvars, eqns, effects={})

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[Union[Var, Literal]]) – list of output variables.

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

• effects (Set[Hashable]) – set of effects. The effects on a jaxpr are a superset of the union of the effects for each equation.

Methods

__init__(constvars, invars, outvars, eqns[, ...])	type constvars Sequence[Var]
pretty_print(*[, source_info, print_shapes, ...])
replace(*[, constvars, invars, outvars, ...])

Attributes

constvars
invars
outvars
eqns
effects