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)
```
This commit is contained in:
Patrick Kidger
2023-11-07 11:34:40 -08:00
parent 9646eff7e1
commit 260fb36876
3 changed files with 82 additions and 32 deletions
+50
View File
@@ -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"]