Bind default arguments (#208)

patrick-kidger · May 19, 2024 · 49cf97e · 49cf97e
1 parent 1ae675f
commit 49cf97e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py
@@ -342,6 +342,7 @@ def modify_annotation(ann):
  @ft.wraps(fn)
  def wrapped_fn(*args, **kwargs): # pyright: ignore
  bound = signature.bind(*args, **kwargs)
+ bound.apply_defaults()
  memos = push_shape_memo(bound.arguments)
  try:
  return fn(*args, **kwargs)
@@ -410,6 +411,7 @@ def wrapped_fn(*args, **kwargs):
  # Raise bind-time errors before we do any shape analysis. (I.e. skip
  # the pointless jaxtyping information for a non-typechecking failure.)
  bound = param_signature.bind(*args, **kwargs)
+ bound.apply_defaults()
 
  memos = push_shape_memo(bound.arguments)
  try:

diff --git a/test/test_decorator.py b/test/test_decorator.py
@@ -117,6 +117,17 @@ def f(x: int, y=1):
  f(1)
 
 
+def test_default_bindings(getkey, jaxtyp, typecheck):
+ @jaxtyp(typecheck)
+ def f(x: int, y: int = 1) -> Float[Array, "x {y}"]:
+ return jr.normal(getkey(), (x, y))
+
+ f(1)
+ f(1, 1)
+ f(1, 0)
+ f(1, 5)
+
+
 class _GlobalFoo:
  pass