From 260fb368760760391b89259a5233f9a1a4fb10ce Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Sat, 21 Oct 2023 12:56:16 -0700 Subject: [PATCH] Fixed mixing variadic+broadcast with variadic+nonbroadcast dimensions. Previously, something like this would not raise an error, as variadic+broadcast dimensions were stored in a separate namespace to variadic+nonbroadcast dimensions: ```python def f(x: Float[Array, "*foo"], y: Float[Array, "#*foo"]): pass a, b = ... assert a.shape == (3, 4) assert b.shape == (5,) f(a, b) ``` --- jaxtyping/_array_types.py | 62 +++++++++++++++++++-------------------- jaxtyping/_decorator.py | 2 +- test/test_array.py | 50 +++++++++++++++++++++++++++++++ 3 files changed, 82 insertions(+), 32 deletions(-) diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py index 06708fc..29befcf 100644 --- a/jaxtyping/_array_types.py +++ b/jaxtyping/_array_types.py @@ -105,7 +105,7 @@ _AbstractDim = Union[Literal[_anonymous_dim], _NamedDim, _FixedDim, _SymbolicDim def _check_dims( cls_dims: list[_AbstractDim], - obj_shape: tuple[int], + obj_shape: tuple[int, ...], single_memo: dict[str, int], ) -> bool: assert len(cls_dims) == len(obj_shape) @@ -201,27 +201,21 @@ class _MetaAbstractArray(type): no_temp_memo = hasattr(storage, "memo_stack") and len(storage.memo_stack) != 0 if no_temp_memo: - single_memo, variadic_memo, variadic_broadcast_memo = storage.memo_stack[-1] + single_memo, variadic_memo = storage.memo_stack[-1] # Make a copy so we don't mutate the original memo during the shape check. single_memo = single_memo.copy() variadic_memo = variadic_memo.copy() - variadic_broadcast_memo = variadic_broadcast_memo.copy() else: # `isinstance` happening outside any @jaxtyped decorators, e.g. at the # global scope. In this case just create a temporary memo, since we're not # going to be comparing against any stored values anyway. single_memo = {} variadic_memo = {} - variadic_broadcast_memo = {} - if cls._check_shape(obj, single_memo, variadic_memo, variadic_broadcast_memo): + if cls._check_shape(obj, single_memo, variadic_memo): # We update the memo every time we successfully pass a shape check if no_temp_memo: - storage.memo_stack[-1] = ( - single_memo, - variadic_memo, - variadic_broadcast_memo, - ) + storage.memo_stack[-1] = single_memo, variadic_memo return True else: return False @@ -230,8 +224,7 @@ class _MetaAbstractArray(type): cls, obj, single_memo: dict[str, int], - variadic_memo: dict[str, tuple[int, ...]], - variadic_broadcast_memo: dict[str, list[tuple[int, ...]]], + variadic_memo: dict[str, tuple[bool, tuple[int, ...]]], ): if cls.index_variadic is None: if obj.ndim != len(cls.dims): @@ -255,30 +248,37 @@ class _MetaAbstractArray(type): return True else: assert type(variadic_dim) is _NamedVariadicDim - variadic_name = variadic_dim.name + name = variadic_dim.name + broadcastable = variadic_dim.broadcastable try: - if variadic_dim.broadcastable: - variadic_shapes = variadic_broadcast_memo[variadic_name] - else: - variadic_shape = variadic_memo[variadic_name] + prev_broadcastable, prev_shape = variadic_memo[name] except KeyError: - if variadic_dim.broadcastable: - variadic_broadcast_memo[variadic_name] = [obj.shape[i:j]] - else: - variadic_memo[variadic_name] = obj.shape[i:j] + variadic_memo[name] = (broadcastable, obj.shape[i:j]) return True else: - if variadic_dim.broadcastable: - new_shape = obj.shape[i:j] - for existing_shape in variadic_shapes: - try: - np.broadcast_shapes(new_shape, existing_shape) - except ValueError: - return False - variadic_shapes.append(new_shape) - return True + new_shape = obj.shape[i:j] + if prev_broadcastable: + try: + broadcast_shape = np.broadcast_shapes(new_shape, prev_shape) + except ValueError: # not broadcastable e.g. (3, 4) and (5,) + return False + if not broadcastable and broadcast_shape != new_shape: + return False + variadic_memo[name] = (broadcastable, broadcast_shape) else: - return variadic_shape == obj.shape[i:j] + if broadcastable: + try: + broadcast_shape = np.broadcast_shapes( + new_shape, prev_shape + ) + except ValueError: # not broadcastable e.g. (3, 4) and (5,) + return False + if broadcast_shape != prev_shape: + return False + else: + if new_shape != prev_shape: + return False + return True assert False diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py index 938aa0e..ca92643 100644 --- a/jaxtyping/_decorator.py +++ b/jaxtyping/_decorator.py @@ -126,7 +126,7 @@ def jaxtyped(fn): memo_stack = storage.memo_stack except AttributeError: memo_stack = storage.memo_stack = [] - memo_stack.append(({}, {}, {})) + memo_stack.append(({}, {})) try: return fn(*args, **kwargs) finally: diff --git a/test/test_array.py b/test/test_array.py index c52a53e..77798cf 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -30,6 +30,7 @@ from jaxtyping import ( AbstractDtype, Array, ArrayLike, + Bool, Float, Float32, jaxtyped, @@ -372,6 +373,55 @@ def test_broadcast_variadic_named(typecheck, getkey): g(o, a) +def test_variadic_mixed_broadcast1(typecheck, getkey): + @jaxtyped + @typecheck + def f(x: Float[Array, " *foo"], y: Float[Array, " #*foo"]): + pass + + a = jr.normal(getkey(), (3, 4)) + b = jr.normal(getkey(), (5,)) + with pytest.raises(ParamError): + f(a, b) + + c = jr.normal(getkey(), (7, 3, 2)) + d = jr.normal(getkey(), (1, 2)) + f(c, d) + + +def test_variadic_mixed_broadcast2(typecheck, getkey): + @jaxtyped + @typecheck + def f(x: Float[Array, " *#foo"], y: Float[Array, " *foo"]): + pass + + a = jr.normal(getkey(), (3, 4)) + b = jr.normal(getkey(), (5,)) + with pytest.raises(ParamError): + f(a, b) + + c = jr.normal(getkey(), (1, 2)) + d = jr.normal(getkey(), (7, 3, 2)) + f(c, d) + + +def test_variadic_mixed_broadcast3(typecheck, getkey): + @jaxtyped + @typecheck + def f( + x: Float[Array, "*B L D"], + *, + y: Float[Array, "*#B J d"], + z: Bool[Array, "*B L J"], + ) -> Float[Array, "*B L D"]: + return x + + x = jr.normal(getkey(), (2, 7, 3, 2, 2)) + y = jr.bernoulli(getkey(), shape=(2, 7, 3, 2, 2)) + z = jr.normal(getkey(), (2, 7, 1, 2, 2)) + f(x, y=z, z=y) + + def test_no_commas(): with pytest.raises(ValueError): Float32[Array, "foo, bar"]