canonical_alias_postprocessor

Handlers for certain “well-known” objects, such as module functions.

Taking the repr of a function or callable usually produces something like

<function vmap at 0x7f98bf556440>

and in some cases produces something like

<jax.custom_derivatives.custom_jvp object at 0x7f98c0b5f130>

This can make it hard to determine what object this actually is from a user perspective, and it would likely be more user-friendly to just output jax.vmap or jax.nn.relu as the representation of this object.

See canonical_aliases for details of how we determine aliases for a subset of common functions.

Functions

replace_with_canonical_aliases(node, path, ...)

Rewrites objects to use well-known aliases when known.