Added print_bindings.

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.

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,

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()))

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"
+    )