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.



An environment that defines a set of canonical aliases.


Expected path where we can find a particular object in a local scope.


Expected path where we can find a particular object in a module.


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.


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.


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.


Checks for newly-imported modules and defines aliases for them.