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

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.

dtype

shape

Methods

evaluate

View source

match

View source

tree_children

View source

tree_map

View source