Thanks for tuning in to Google I/O. View all sessions on demandWatch on demand

oryx.experimental.matching.jax_rewrite.Primitive

A JAXExpression corresponding to a jax.core.Primitive.

Inherits From: JaxExpression, Expression, Pattern

oryx.experimental.matching.jax_rewrite.Primitive(
    primitive: Union[oryx.experimental.matching.matcher.Pattern, jax_core.Primitive],
    operands: Tuple[Any, ...],
    params: Union[oryx.experimental.matching.matcher.Pattern, oryx.experimental.matching.jax_rewrite.Params]
)

A Primitive encapsulates a jax.core.Primitive, its arguments and its parameters.

primitive A JAX primitive.
operands A tuple of expressions that are evaluated and passed into the primitive when a Primitive is evaluated.
params A Params object that holds the parameters passed into the JAX primitive when a Primitive is evaluated.
dtype

shape

shape_dtype

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
) -> 'Primitive'

__eq__

View source

__eq__(
    other: Any
) -> bool

Return self==value.