oryx.experimental.matching.jax_rewrite.JaxVar

View source on GitHub

An expression that looks up a provided name in the environment.

Inherits From: JaxExpression, Expression, Pattern

oryx.experimental.matching.jax_rewrite.JaxVar(
    name: str, shape: Tuple[Any], dtype: jnp.dtype
)

Args
name	a string name corresponding the key used to lookup this JaxVar's value in an environment.
shape	a tuple of integers corresponding to the shape of the JaxVar.
dtype	a JAX dtype object for this JaxVar's dtype.

Attributes
name	The string name for the JaxVar used to look up the JaxVar's value in an environment.
dtype
shape

Methods

evaluate

View source

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() -> Iterator[Expr]

tree_map

View source

tree_map(
    _
) -> 'JaxVar'