Added with jaxtyped("context"):, for now undocumented.

jaxtyping/_decorator.py

@@ -118,6 +118,8 @@ def jaxtyped(fn):
         else:
             fdel = jaxtyped(fn.fdel)
         return property(fget=fget, fset=fset, fdel=fdel)
+    elif fn == "context":
+        return _JaxtypingContext()
     else:
 
         @ft.wraps(fn)
@@ -136,6 +138,18 @@ def jaxtyped(fn):
         return wrapped_fn
 
 
+class _JaxtypingContext:
+    def __enter__(self):
+        try:
+            memo_stack = storage.memo_stack
+        except AttributeError:
+            memo_stack = storage.memo_stack = []
+        memo_stack.append(({}, {}))
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        storage.memo_stack.pop()
+
+
 @jaxtyped
 def _check_dataclass_annotations(self, typechecker):
     for field in dataclasses.fields(self):

test/test_decorator.py

@@ -1,6 +1,8 @@
 import abc
 
-from jaxtyping import jaxtyped
+import jax.random as jr
+
+from jaxtyping import Array, Float, jaxtyped
 
 
 class M(metaclass=abc.ABCMeta):
@@ -75,3 +77,11 @@ def test_abstractmethod():
 def test_property():
     assert N().j1 == 3
     assert N().j2 == 4
+
+
+def test_context(getkey):
+    a = jr.normal(getkey(), (3, 4))
+    b = jr.normal(getkey(), (5,))
+    with jaxtyped("context"):
+        assert isinstance(a, Float[Array, "foo bar"])
+        assert not isinstance(b, Float[Array, "foo"])