jax.extend.core.Jaxpr#

class jax.extend.core.Jaxpr(constvars, invars, outvars, eqns, effects=frozenset({}), debug_info=None)[源代码]#
参数:

  • constvars (Sequence[Var])

  • invars (Sequence[Var])

  • outvars (Sequence[Atom])

  • eqns (Sequence[JaxprEqn])

  • effects (Effects)

  • debug_info (JaxprDebugInfo | None)

__init__(constvars, invars, outvars, eqns, effects=frozenset({}), debug_info=None)[源代码]#
参数:

  • constvars (Sequence[Var]) – 为常量引入的变量列表。数组常量被替换为这样的变量，而标量常量则保持内联。

  • invars (Sequence[Var]) – 输入变量列表。constvars 和 invars 一起是 Jaxpr 的输入。

  • outvars (Sequence[Atom]) – 输出原子列表。

  • eqns (Sequence[JaxprEqn]) – 方程列表。

  • effects (Effects) – 效果集合。jaxpr 上的效果是每个方程效果的并集的超集。

  • debug_info (JaxprDebugInfo | None | None) – 可选的 JaxprDebugInfo。

方法

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

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

replace(**kwargs)

属性

constvars

debug_info

effects

eqns

invars

outvars