oryx.experimental.matching.jax_rewrite.BoundExpression

View source on GitHub

Represents JAX expressions with closed over constants.

Inherits From: JaxExpression, Expression, Pattern

oryx.experimental.matching.jax_rewrite.BoundExpression(
    expressions: Sequence[Expr],
    consts: oryx.experimental.matching.jax_rewrite.Bindings
)

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.

Attributes
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

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

Evaluates using an environment augmented with constants.

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

__eq__

__eq__(
    other
)