oryx.experimental.matching.jax_rewrite.JaxExpression

View source on GitHub

A node in an expression tree.

Inherits From: Expression, Pattern

JAXExpressions are subclasses of rules.Expression so if the tree_map and tree_children methods are implemented, it is compatible with the rewrite interface in rules. To be compatible with the pattern matching interface in matcher, the match method needs to be implemented too.

Attributes
dtype
shape

Methods

evaluate

View source

@abc.abstractmethod
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()

tree_map

View source

tree_map(
    fn
)