oryx.experimental.matching.jax_rewrite.CallPrimitive

Encapsulates JAX CallPrimitives like jit and pmap.

Inherits From: JaxExpression, Expression, Pattern

oryx.experimental.matching.jax_rewrite.CallPrimitive(
    primitive: jax_core.Primitive,
    operands: Sequence[Any],
    expression: Any,
    params: oryx.experimental.matching.jax_rewrite.Params,
    variable_names: Sequence[str]
)

primitive A JAX call primitive.
operands A sequence of expressions that are evaluated and passed as inputs to the primitive when the CallPrimitive is evaluated.
expression The expression that corresponds to the body of the call primitive. The operands are bound to the variable_names in an environment and the expression is evaluated with that environment.
params A Params object corresponding to the parameters of the call primitive.
variable_names A sequence of string names that are used as keys for the operands in the environment expression is evaluated in.
dtype

shape

Methods

evaluate

View source

evaluate(
    env: oryx.experimental.matching.jax_rewrite.Bindings
) -> Any

match

View source

match(
    expr: Expr,
    bindings: oryx.experimental.matching.jax_rewrite.Bindings,
    succeed: oryx.experimental.matching.jax_rewrite.Continuation
) -> oryx.experimental.matching.jax_rewrite.Success

tree_children

View source

tree_children() -> Iterator[Expr]

tree_map

View source

tree_map(
    fn
) -> 'CallPrimitive'

__eq__

__eq__(
    other
)