treescope.canonical_aliases

treescope.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 and Equinox both 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`(include[, excludes])	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.