mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-09-09 11:24:55 +08:00
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:
@@ -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"]
|
||||
|
||||
Reference in New Issue
Block a user