jax.core.Jaxpr

jax.core.Jaxpr#

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

  • 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)[source]#
参数：:

  • 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

目录