ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

oryx.experimental.matching.jax_rewrite.Primitive

A JAXExpression corresponding to a jax.core.Primitive.

Inherits From: JaxExpression, Expression, Pattern

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

match

View source

tree_children

View source

tree_map

View source

__eq__

View source

Return self==value.