jax.extend.core.ClosedJaxpr#

class jax.extend.core.ClosedJaxpr(jaxpr, consts)[source]#
Parameters:

  • jaxpr (Jaxpr)

  • consts (Sequence)

__init__(jaxpr, consts)[source]#
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

final_aval_qdds

in_aval_qdds

in_avals

jaxpr

literals

out_avals