mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-08-21 10:25:22 +08:00
Roman Knyazhitskiy
172b83b4fc
* Add a test for generators * Remove output annotations from decorators Also guarded torch imports for better compatibility with requirements.txt * Add flag to the main meta class to skip the typecheck * Return to the old solution * Make async tests work * Minor adjustments/fixing typos * Correct Python path for new tests * Remove some jax-dependent code * Implement equality for MetaArrays * Make all Dim variations frozen dataclasses * Shorten AbstractArray methods * Final touches * Removing get_origin use * Update tests with @jaxtyp
695 lines
17 KiB
Python
695 lines
17 KiB
Python
# Copyright (c) 2022 Google LLC
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy of
|
|
# this software and associated documentation files (the "Software"), to deal in
|
|
# the Software without restriction, including without limitation the rights to
|
|
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
|
|
# the Software, and to permit persons to whom the Software is furnished to do so,
|
|
# subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in all
|
|
# copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
|
|
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
|
|
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
|
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
|
|
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
|
|
import dataclasses as dc
|
|
import sys
|
|
from typing import get_args, get_origin, Union
|
|
|
|
import jax.numpy as jnp
|
|
import jax.random as jr
|
|
import numpy as np
|
|
import pytest
|
|
|
|
|
|
try:
|
|
import torch
|
|
except ImportError:
|
|
torch = None
|
|
|
|
from jaxtyping import (
|
|
AbstractDtype,
|
|
AnnotationError,
|
|
Array,
|
|
ArrayLike,
|
|
Bool,
|
|
Float,
|
|
Float32,
|
|
jaxtyped,
|
|
Key,
|
|
PRNGKeyArray,
|
|
Scalar,
|
|
Shaped,
|
|
)
|
|
|
|
from .helpers import ParamError, ReturnError
|
|
|
|
|
|
def test_basic(jaxtyp, typecheck):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Shaped[Array, "..."]):
|
|
pass
|
|
|
|
g(jnp.array(1.0))
|
|
|
|
|
|
def test_dtypes():
|
|
from jaxtyping import ( # noqa: F401
|
|
Array,
|
|
BFloat16,
|
|
Bool,
|
|
Complex,
|
|
Complex64,
|
|
Complex128,
|
|
Float,
|
|
Float16,
|
|
Float32,
|
|
Float64,
|
|
Inexact,
|
|
Int,
|
|
Int4,
|
|
Int8,
|
|
Int16,
|
|
Int32,
|
|
Int64,
|
|
Num,
|
|
Shaped,
|
|
UInt,
|
|
UInt4,
|
|
UInt8,
|
|
UInt16,
|
|
UInt32,
|
|
UInt64,
|
|
)
|
|
|
|
for key, val in locals().items():
|
|
if issubclass(val, AbstractDtype):
|
|
assert key == val.__name__
|
|
|
|
|
|
def test_return(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Float[Array, "b c"]) -> Float[Array, "c b"]:
|
|
return jnp.transpose(x)
|
|
|
|
g(jr.normal(getkey(), (3, 4)))
|
|
|
|
@jaxtyp(typecheck)
|
|
def h(x: Float[Array, "b c"]) -> Float[Array, "b c"]:
|
|
return jnp.transpose(x)
|
|
|
|
with pytest.raises(ReturnError):
|
|
h(jr.normal(getkey(), (3, 4)))
|
|
|
|
|
|
def test_two_args(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]):
|
|
return x @ y
|
|
|
|
g(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (4, 5)))
|
|
with pytest.raises(ParamError):
|
|
g(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4)))
|
|
|
|
@jaxtyp(typecheck)
|
|
def h(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]) -> Shaped[Array, "b d"]:
|
|
return x @ y
|
|
|
|
h(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (4, 5)))
|
|
with pytest.raises(ParamError):
|
|
h(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4)))
|
|
|
|
|
|
def test_any_dtype(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Shaped[Array, "a b"]) -> Shaped[Array, "a b"]:
|
|
return x
|
|
|
|
g(jr.normal(getkey(), (3, 4)))
|
|
g(jnp.array([[True, False]]))
|
|
g(jnp.array([[1, 2], [3, 4]], dtype=jnp.int4))
|
|
g(jnp.array([[1, 2], [3, 4]], dtype=jnp.int8))
|
|
g(jnp.array([[1, 2], [3, 4]], dtype=jnp.uint4))
|
|
g(jnp.array([[1, 2], [3, 4]], dtype=jnp.uint16))
|
|
g(jr.normal(getkey(), (3, 4), dtype=jnp.complex128))
|
|
g(jr.normal(getkey(), (3, 4), dtype=jnp.bfloat16))
|
|
|
|
with pytest.raises(ParamError):
|
|
g(jr.normal(getkey(), (1,)))
|
|
|
|
|
|
def test_nested_jaxtyped(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Float32[Array, "b c"], transpose: bool) -> Float32[Array, "c b"]:
|
|
return h(x, transpose)
|
|
|
|
@jaxtyp(typecheck)
|
|
def h(x: Float32[Array, "c b"], transpose: bool) -> Float32[Array, "b c"]:
|
|
if transpose:
|
|
return jnp.transpose(x)
|
|
else:
|
|
return x
|
|
|
|
g(jr.normal(getkey(), (2, 3)), True)
|
|
g(jr.normal(getkey(), (3, 3)), True)
|
|
g(jr.normal(getkey(), (3, 3)), False)
|
|
with pytest.raises(ReturnError):
|
|
g(jr.normal(getkey(), (2, 3)), False)
|
|
|
|
|
|
def test_nested_nojaxtyped(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Float32[Array, "b c"]):
|
|
return h(x)
|
|
|
|
@typecheck
|
|
def h(x: Float32[Array, "c b"]):
|
|
return x
|
|
|
|
with pytest.raises(ParamError):
|
|
g(jr.normal(getkey(), (2, 3)))
|
|
|
|
|
|
def test_isinstance(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Float32[Array, "b c"]) -> Float32[Array, " z"]:
|
|
y = jnp.transpose(x)
|
|
assert isinstance(y, Float32[Array, "c b"])
|
|
assert not isinstance(
|
|
y, Float32[Array, "b z"]
|
|
) # z left unbound as b!=c (unless x symmetric, which it isn't)
|
|
out = jr.normal(getkey(), (500,))
|
|
assert isinstance(out, Float32[Array, "z"]) # z now bound
|
|
return out
|
|
|
|
g(jr.normal(getkey(), (2, 3)))
|
|
|
|
|
|
def test_fixed(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(
|
|
x: Float32[Array, "4 5 foo"], y: Float32[Array, " foo"]
|
|
) -> Float32[Array, "4 5"]:
|
|
return x @ y
|
|
|
|
a = jr.normal(getkey(), (4, 5, 2))
|
|
b = jr.normal(getkey(), (2,))
|
|
assert g(a, b).shape == (4, 5)
|
|
|
|
c = jr.normal(getkey(), (3, 5, 2))
|
|
with pytest.raises(ParamError):
|
|
g(c, b)
|
|
|
|
|
|
def test_anonymous(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Float32[Array, "foo _"], y: Float32[Array, " _"]):
|
|
pass
|
|
|
|
a = jr.normal(getkey(), (3, 4))
|
|
b = jr.normal(getkey(), (5,))
|
|
g(a, b)
|
|
|
|
|
|
def test_named_variadic(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(
|
|
x: Float32[Array, "*batch foo"],
|
|
y: Float32[Array, " *batch"],
|
|
z: Float32[Array, " foo"],
|
|
):
|
|
pass
|
|
|
|
c = jr.normal(getkey(), (5,))
|
|
|
|
a1 = jr.normal(getkey(), (5,))
|
|
b1 = jr.normal(getkey(), ())
|
|
g(a1, b1, c)
|
|
|
|
a2 = jr.normal(getkey(), (3, 5))
|
|
b2 = jr.normal(getkey(), (3,))
|
|
g(a2, b2, c)
|
|
|
|
with pytest.raises(ParamError):
|
|
g(a1, b2, c)
|
|
with pytest.raises(ParamError):
|
|
g(a2, b1, c)
|
|
|
|
@jaxtyp(typecheck)
|
|
def h(x: Float32[Array, " foo *batch"], y: Float32[Array, " foo *batch bar"]):
|
|
pass
|
|
|
|
a = jr.normal(getkey(), (4,))
|
|
b = jr.normal(getkey(), (4, 3))
|
|
c = jr.normal(getkey(), (3, 4))
|
|
h(a, b)
|
|
with pytest.raises(ParamError):
|
|
h(a, c)
|
|
with pytest.raises(ParamError):
|
|
h(b, c)
|
|
|
|
|
|
def test_anonymous_variadic(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Float32[Array, "... foo"], y: Float32[Array, " foo"]):
|
|
pass
|
|
|
|
a1 = jr.normal(getkey(), (5,))
|
|
a2 = jr.normal(getkey(), (3, 5))
|
|
a3 = jr.normal(getkey(), (3, 4, 5))
|
|
b = jr.normal(getkey(), (5,))
|
|
c = jr.normal(getkey(), (1,))
|
|
g(a1, b)
|
|
g(a2, b)
|
|
g(a3, b)
|
|
with pytest.raises(ParamError):
|
|
g(a1, c)
|
|
with pytest.raises(ParamError):
|
|
g(a2, c)
|
|
with pytest.raises(ParamError):
|
|
g(a3, c)
|
|
|
|
|
|
def test_broadcast_fixed(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Float32[Array, "#4"]):
|
|
pass
|
|
|
|
g(jr.normal(getkey(), (4,)))
|
|
g(jr.normal(getkey(), (1,)))
|
|
|
|
with pytest.raises(ParamError):
|
|
g(jr.normal(getkey(), (3,)))
|
|
|
|
|
|
def test_broadcast_named(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Float32[Array, " #foo"], y: Float32[Array, " #foo"]):
|
|
pass
|
|
|
|
a = jr.normal(getkey(), (3,))
|
|
b = jr.normal(getkey(), (4,))
|
|
c = jr.normal(getkey(), (1,))
|
|
|
|
g(a, a)
|
|
g(b, b)
|
|
g(c, c)
|
|
g(a, c)
|
|
g(b, c)
|
|
g(c, a)
|
|
g(c, b)
|
|
|
|
with pytest.raises(ParamError):
|
|
g(a, b)
|
|
with pytest.raises(ParamError):
|
|
g(b, a)
|
|
|
|
|
|
def test_broadcast_variadic_named(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def g(x: Float32[Array, " *#foo"], y: Float32[Array, " *#foo"]):
|
|
pass
|
|
|
|
a = jr.normal(getkey(), (3,))
|
|
b = jr.normal(getkey(), (4,))
|
|
c = jr.normal(getkey(), (4, 4))
|
|
d = jr.normal(getkey(), (5, 6))
|
|
|
|
j = jr.normal(getkey(), (1,))
|
|
k = jr.normal(getkey(), (1, 4))
|
|
l = jr.normal(getkey(), (5, 1)) # noqa: E741
|
|
m = jr.normal(getkey(), (1, 1))
|
|
n = jr.normal(getkey(), (2, 1))
|
|
o = jr.normal(getkey(), (1, 6))
|
|
|
|
g(a, a)
|
|
g(b, b)
|
|
g(c, c)
|
|
g(d, d)
|
|
g(b, c)
|
|
with pytest.raises(ParamError):
|
|
g(a, b)
|
|
with pytest.raises(ParamError):
|
|
g(a, c)
|
|
with pytest.raises(ParamError):
|
|
g(a, b)
|
|
with pytest.raises(ParamError):
|
|
g(d, b)
|
|
|
|
g(a, j)
|
|
g(b, j)
|
|
g(c, j)
|
|
g(d, j)
|
|
g(b, k)
|
|
g(c, k)
|
|
with pytest.raises(ParamError):
|
|
g(d, k)
|
|
with pytest.raises(ParamError):
|
|
g(c, l)
|
|
g(d, l)
|
|
g(a, m)
|
|
g(c, m)
|
|
g(d, m)
|
|
g(a, n)
|
|
g(b, n)
|
|
with pytest.raises(ParamError):
|
|
g(c, n)
|
|
with pytest.raises(ParamError):
|
|
g(d, n)
|
|
g(o, d)
|
|
with pytest.raises(ParamError):
|
|
g(o, c)
|
|
with pytest.raises(ParamError):
|
|
g(o, a)
|
|
|
|
|
|
def test_variadic_mixed_broadcast(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(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(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(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(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(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"]
|
|
|
|
|
|
def test_symbolic(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
|
|
return x[1:]
|
|
|
|
@jaxtyp(typecheck)
|
|
def cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
|
|
return jnp.concatenate([x, x])
|
|
|
|
@jaxtyp(typecheck)
|
|
def bad_make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
|
|
return x
|
|
|
|
@jaxtyp(typecheck)
|
|
def bad_cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
|
|
return jnp.concatenate([x, x, x])
|
|
|
|
x = jr.normal(getkey(), (5,))
|
|
assert make_slice(x).shape == (4,)
|
|
assert cat(x).shape == (10,)
|
|
|
|
y = jr.normal(getkey(), (3, 4))
|
|
with pytest.raises(ParamError):
|
|
make_slice(y)
|
|
with pytest.raises(ParamError):
|
|
cat(y)
|
|
|
|
with pytest.raises(ReturnError):
|
|
bad_make_slice(x)
|
|
with pytest.raises(ReturnError):
|
|
bad_cat(x)
|
|
|
|
|
|
def test_incomplete_symbolic(jaxtyp, typecheck, getkey):
|
|
@jaxtyp(typecheck)
|
|
def foo(x: Float32[Array, " 2*dim"]):
|
|
pass
|
|
|
|
x = jr.normal(getkey(), (4,))
|
|
with pytest.raises(AnnotationError):
|
|
foo(x)
|
|
|
|
|
|
def test_deferred_symbolic_good(jaxtyp, typecheck):
|
|
@jaxtyp(typecheck)
|
|
def foo(dim: int, fill: Float[Array, ""]) -> Float[Array, " {dim}"]:
|
|
return jnp.full((dim,), fill)
|
|
|
|
class A:
|
|
size = 5
|
|
|
|
@jaxtyp(typecheck)
|
|
def bar(self, fill: Float[Array, ""]) -> Float[Array, " {self.size}"]:
|
|
return jnp.full((self.size,), fill)
|
|
|
|
foo(3, jnp.array(0.0))
|
|
A().bar(jnp.array(0.0))
|
|
|
|
|
|
def test_deferred_symbolic_bad(jaxtyp, typecheck):
|
|
@jaxtyp(typecheck)
|
|
def foo(dim: int, fill: Float[Array, ""]) -> Float[Array, " {dim-1}"]:
|
|
return jnp.full((dim,), fill)
|
|
|
|
class A:
|
|
size = 5
|
|
|
|
@jaxtyp(typecheck)
|
|
def bar(self, fill: Float[Array, ""]) -> Float[Array, " {self.size}-1"]:
|
|
return jnp.full((self.size,), fill)
|
|
|
|
with pytest.raises(ReturnError):
|
|
foo(3, jnp.array(0.0))
|
|
|
|
with pytest.raises(ReturnError):
|
|
A().bar(jnp.array(0.0))
|
|
|
|
|
|
def test_deferred_symbolic_dataclass(typecheck):
|
|
@jaxtyped(typechecker=typecheck)
|
|
@dc.dataclass
|
|
class A:
|
|
value: int
|
|
array: Float[Array, " {self.value}"]
|
|
|
|
A(3, jnp.zeros(3))
|
|
|
|
with pytest.raises(ParamError):
|
|
A(3, jnp.zeros(4))
|
|
|
|
|
|
def test_arraylike(typecheck, getkey):
|
|
floatlike1 = Float32[ArrayLike, ""]
|
|
floatlike2 = Float[ArrayLike, ""]
|
|
floatlike3 = Float32[ArrayLike, "4"]
|
|
|
|
assert get_origin(floatlike1) is Union
|
|
assert get_origin(floatlike2) is Union
|
|
assert get_origin(floatlike3) is Union
|
|
assert set(get_args(floatlike1)) == {
|
|
Float32[Array, ""],
|
|
Float32[np.ndarray, ""],
|
|
Float32[np.number, ""],
|
|
float,
|
|
}
|
|
assert set(get_args(floatlike2)) == {
|
|
Float[Array, ""],
|
|
Float[np.ndarray, ""],
|
|
Float[np.number, ""],
|
|
float,
|
|
}
|
|
assert set(get_args(floatlike3)) == {
|
|
Float32[Array, "4"],
|
|
Float32[np.ndarray, "4"],
|
|
}
|
|
|
|
shaped1 = Shaped[ArrayLike, ""]
|
|
shaped2 = Shaped[ArrayLike, "4"]
|
|
assert get_origin(shaped1) is Union
|
|
assert get_origin(shaped2) is Union
|
|
assert set(get_args(shaped1)) == {
|
|
Shaped[Array, ""],
|
|
Shaped[np.ndarray, ""],
|
|
Shaped[np.bool_, ""],
|
|
Shaped[np.number, ""],
|
|
bool,
|
|
int,
|
|
float,
|
|
complex,
|
|
}
|
|
assert set(get_args(shaped2)) == {
|
|
Shaped[Array, "4"],
|
|
Shaped[np.ndarray, "4"],
|
|
}
|
|
|
|
|
|
def test_subclass():
|
|
assert issubclass(Float[Array, ""], Array)
|
|
assert issubclass(Float[np.ndarray, ""], np.ndarray)
|
|
|
|
if torch is not None:
|
|
assert issubclass(Float[torch.Tensor, ""], torch.Tensor)
|
|
|
|
|
|
def test_ignored_names():
|
|
x = Float[np.ndarray, "foo=4"]
|
|
|
|
assert isinstance(np.zeros(4), x)
|
|
assert not isinstance(np.zeros(5), x)
|
|
assert not isinstance(np.zeros((4, 5)), x)
|
|
|
|
y = Float[np.ndarray, "bar qux foo=bar+qux"]
|
|
|
|
assert isinstance(np.zeros((2, 3, 5)), y)
|
|
assert not isinstance(np.zeros((2, 3, 6)), y)
|
|
|
|
z = Float[np.ndarray, "bar #foo=bar"]
|
|
|
|
assert isinstance(np.zeros((3, 3)), z)
|
|
assert isinstance(np.zeros((3, 1)), z)
|
|
assert not isinstance(np.zeros((3, 4)), z)
|
|
|
|
# Weird but legal
|
|
w = Float[np.ndarray, "bar foo=#bar"]
|
|
|
|
assert isinstance(np.zeros((3, 3)), w)
|
|
assert isinstance(np.zeros((3, 1)), w)
|
|
assert not isinstance(np.zeros((3, 4)), w)
|
|
|
|
|
|
def test_symbolic_functions():
|
|
x = Float[np.ndarray, "foo bar min(foo,bar)"]
|
|
|
|
assert isinstance(np.zeros((2, 3, 2)), x)
|
|
assert isinstance(np.zeros((3, 2, 2)), x)
|
|
assert not isinstance(np.zeros((3, 2, 4)), x)
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires Python 3.10")
|
|
def test_py310_unions():
|
|
x = np.zeros(3)
|
|
y = Shaped[Array | np.ndarray, "_"]
|
|
assert isinstance(x, get_args(y))
|
|
|
|
|
|
def test_key(jaxtyp, typecheck):
|
|
@jaxtyp(typecheck)
|
|
def f(x: PRNGKeyArray):
|
|
pass
|
|
|
|
f(jr.key(0))
|
|
f(jr.PRNGKey(0))
|
|
|
|
with pytest.raises(ParamError):
|
|
f(object())
|
|
with pytest.raises(ParamError):
|
|
f(1)
|
|
with pytest.raises(ParamError):
|
|
f(jnp.array(3))
|
|
with pytest.raises(ParamError):
|
|
f(jnp.array(3.0))
|
|
|
|
|
|
def test_key_dtype(jaxtyp, typecheck):
|
|
@jaxtyp(typecheck)
|
|
def f1(x: Key[Array, ""]):
|
|
pass
|
|
|
|
@jaxtyp(typecheck)
|
|
def f2(x: Key[Scalar, ""]):
|
|
pass
|
|
|
|
for f in (f1, f2):
|
|
f(jr.key(0))
|
|
|
|
with pytest.raises(ParamError):
|
|
f(jr.PRNGKey(0))
|
|
with pytest.raises(ParamError):
|
|
f(object())
|
|
with pytest.raises(ParamError):
|
|
f(1)
|
|
with pytest.raises(ParamError):
|
|
f(jnp.array(3))
|
|
with pytest.raises(ParamError):
|
|
f(jnp.array(3.0))
|
|
|
|
|
|
def test_extension(jaxtyp, typecheck, getkey):
|
|
X = Shaped[Array, "a b"]
|
|
Y = Shaped[X, "c d"]
|
|
Z = Shaped[Array, "c d a b"]
|
|
assert str(Z) == str(Y)
|
|
|
|
X = Float[Array, "a"]
|
|
Y = Float[X, "b"]
|
|
|
|
@jaxtyp(typecheck)
|
|
def f(a: X, b: Y):
|
|
...
|
|
|
|
a = jr.normal(getkey(), (3, 4))
|
|
b = jr.normal(getkey(), (4,))
|
|
c = jr.normal(getkey(), (3,))
|
|
|
|
f(b, a)
|
|
with pytest.raises(ParamError):
|
|
f(c, a)
|
|
with pytest.raises(ParamError):
|
|
f(a, a)
|
|
|
|
@typecheck
|
|
def g(a: Shaped[PRNGKeyArray, "2"]):
|
|
...
|
|
|
|
with pytest.raises(ParamError):
|
|
g(jr.PRNGKey(0))
|
|
g(jr.split(jr.PRNGKey(0)))
|
|
with pytest.raises(ParamError):
|
|
g(jr.split(jr.PRNGKey(0), 3))
|
|
|
|
|
|
def test_scalar_variadic_dim():
|
|
assert Float[float, "..."] is float
|
|
assert Float[float, "#*shape"] is float
|
|
|
|
# This one is a bit weird -- it should really also assert that shape==(), but we
|
|
# don't implement that.
|
|
assert Float[float, "*shape"] is float
|
|
|
|
|
|
def test_scalar_dtype_mismatch():
|
|
with pytest.raises(ValueError):
|
|
Float[bool, "..."]
|