diff --git a/docs/api/advanced-features.md b/docs/api/advanced-features.md index edb07bf..ad679cd 100644 --- a/docs/api/advanced-features.md +++ b/docs/api/advanced-features.md @@ -7,6 +7,10 @@ members: false +## Printing axis bindings + +::: jaxtyping.print_bindings + ## Introspection If you're writing your own type hint parser, then you may wish to detect if some Python object is a jaxtyping-provided type. diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 991f1b1..1bce0ff 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -30,7 +30,7 @@ from ._array_types import ( set_array_name_format as set_array_name_format, ) from ._config import config as config -from ._decorator import jaxtyped as jaxtyped +from ._decorator import jaxtyped as jaxtyped, print_bindings as print_bindings from ._errors import ( AnnotationError as AnnotationError, TypeCheckError as TypeCheckError, diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py index af18be9..f4ef9ac 100644 --- a/jaxtyping/_decorator.py +++ b/jaxtyping/_decorator.py @@ -36,7 +36,7 @@ else: from ._config import config from ._errors import AnnotationError, TypeCheckError -from ._storage import pop_shape_memo, push_shape_memo +from ._storage import get_shape_memo, pop_shape_memo, push_shape_memo class _Sentinel: @@ -808,3 +808,21 @@ def _no_jaxtyping_note(e: Exception) -> bool: _spacer = "--------------------\n" + + +def print_bindings(): + """Prints the values of the current jaxtyping axis bindings. Intended for debugging. + + That is, whilst doing runtime type checking, so that e.g. the `foo` and `bar` of + `Float[Array, "foo bar"]` are assigned values -- this function will print out those + values. + + **Arguments:** + + Nothing. + + **Returns:** + + Nothing. + """ + print(_exc_shape_info(get_shape_memo())) diff --git a/test/test_decorator.py b/test/test_decorator.py index 0ab3836..c6d3c1f 100644 --- a/test/test_decorator.py +++ b/test/test_decorator.py @@ -1,9 +1,10 @@ import abc +import jax.numpy as jnp import jax.random as jr import pytest -from jaxtyping import Array, Float, jaxtyped +from jaxtyping import Array, Float, jaxtyped, print_bindings from .helpers import ParamError, ReturnError @@ -166,3 +167,17 @@ def test_local_stringified_annotation(typecheck): # We don't check that errors are raised if it goes wrong, since we can't usually # resolve local type annotations at runtime. Best we can hope for is not to raise # a spurious error about not being able to find the type. + + +def test_print_bindings(typecheck, capfd): + @jaxtyped(typechecker=typecheck) + def f(x: Float[Array, "foo bar"]): + print_bindings() + + capfd.readouterr() + f(jnp.zeros((3, 4))) + text, _ = capfd.readouterr() + assert text == ( + "The current values for each jaxtyping axis annotation are as follows." + "\nfoo=3\nbar=4\n" + )