oryx.experimental.matching.jax_rewrite.JaxExpression

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`

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
)