canonical_aliases

Registry of 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.

Many functions and classes store a reference to the location where they were originally defined (in their __module__ and __qualname__ attributes), which can be used to find an alias for them. However, this may not be the “canonical” alias, because some modules re-export private symbols under a public namespace (in particular, JAX, Penzai, and Equinox all do this).

This module contains a registry of canonical paths for specific functions and other objects, so that they can be rendered in a predictable way by treescope. The intended purpose is primarily for interactive printing and debugging, although it also helps for reifying objects into executable code through round-trippable pretty printing, which can enable a simple form of serialization.

This module also supports walking the public API of a package to automatically set aliases; this is done by default for JAX and a few other libraries to ensure the pretty-printed outputs avoid private module paths whenever possible. This is intended as a heuristic to construct readable aliases for common objects on a best-effort basis. It is not guaranteed that these inferred aliases will always be stable across different versions of the external libraries.

Classes

CanonicalAliasEnvironment	An environment that defines a set of canonical aliases.
LocalNamePath	Expected path where we can find a particular object in a local scope.
ModuleAttributePath	Expected path where we can find a particular object in a module.

Functions

add_alias(the_object, path[, on_conflict])	Adds an alias to this object.
default_well_known_filter(the_object, path)	Checks if an object looks like something we want to define an alias for.
lookup_alias(the_object[, ...])	Retrieves an alias for this object, if possible.
maybe_local_module_name(module)	Returns a name for this module, possibly looking up local aliases.
populate_from_public_api(module[, predicate])	Populates canonical aliases with all public symbols in a module.
prefix_filter(prefix)	Builds a filter that only defines aliases within a given prefix.
relative_alias_names(relative_scope[, predicate])	Context manager that makes lookup_alias return relative aliases.
update_lazy_aliases()	Checks for newly-imported modules and defines aliases for them.