jax.extend.core.Jaxpr

jax.extend.core.Jaxpr#

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

  • constvars (Sequence[Var] | Jaxpr)

  • invars (Sequence[Var] | Sequence[Any] | None)

  • outvars (Sequence[Atom] | None)

  • eqns (Sequence[JaxprEqn] | None)

  • effects (Effects)

  • debug_info (DebugInfo)

  • is_high (bool)

  • consts (Sequence[Any] | None)

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

  • constvars (Sequence[Var] | Jaxpr)

  • invars (Sequence[Var] | Sequence[Any] | None)

  • outvars (Sequence[Atom] | None)

  • eqns (Sequence[JaxprEqn] | None)

  • effects (Effects)

  • debug_info (DebugInfo)

  • is_high (bool)

  • consts (Sequence[Any] | None)

Methods

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

map_jaxpr(f)

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

replace(**kwargs)

with_consts(consts)

Returns a copy of this jaxpr with consts attached as the values of its constant inputs (see is_closed).

Attributes

all_invars

consts

constvars

debug_info

effects

eqns

final_aval_qdds

in_aval_qdds

in_avals

invars

is_closed

is_high

jaxpr

literals

num_consts

out_avals

outvars