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

oryx.experimental.matching.jax_rewrite.BoundExpression

Represents JAX expressions with closed over constants.

Inherits From: JaxExpression, Expression, Pattern

A BoundExpression enables pinning JaxVars in an expression to fixed values, removing the need to bind them to values when the BoundExpression is evaluated. Conceptually this is equivalent to a jax.core.ClosedJaxpr.

expressions A sequence of expressions that are evaluated to produce the result of this BoundExpression.
consts A dictionary mapping string names (corresponding to JaxVars in expressions) to their JAX values.
dtype

shape

Methods

evaluate

View source

Evaluates using an environment augmented with constants.

match

View source

tree_children

View source

tree_map

View source

__eq__