mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-08-21 11:16:29 +08:00
105 lines
3.2 KiB
Python
105 lines
3.2 KiB
Python
from typing import Any
|
|
|
|
import equinox as eqx
|
|
import jax.numpy as jnp
|
|
import pytest
|
|
|
|
from jaxtyping import Array, Float, jaxtyped, PyTree, TypeCheckError
|
|
|
|
|
|
def test_arg_localisation(typecheck):
|
|
@jaxtyped(typechecker=typecheck)
|
|
def f(x: str, y: str, z: int):
|
|
pass
|
|
|
|
matches = [
|
|
"Type-check error whilst checking the parameters of f",
|
|
"The problem arose whilst typechecking parameter 'z'.",
|
|
"Called with parameters: {'x': 'hi', 'y': 'bye', 'z': 'not-an-int'}",
|
|
r"Parameter annotations: \(x: str, y: str, z: int\).",
|
|
]
|
|
for match in matches:
|
|
with pytest.raises(TypeCheckError, match=match):
|
|
f("hi", "bye", "not-an-int")
|
|
|
|
@jaxtyped(typechecker=typecheck)
|
|
def g(x: Float[Array, "a b"], y: Float[Array, "b c"]):
|
|
pass
|
|
|
|
x = jnp.zeros((2, 3))
|
|
y = jnp.zeros((4, 3))
|
|
matches = [
|
|
"Type-check error whilst checking the parameters of g",
|
|
"The problem arose whilst typechecking parameter 'y'.",
|
|
r"Called with parameters: {'x': f32\[2,3\], 'y': f32\[4,3\]}",
|
|
(
|
|
r"Parameter annotations: \(x: Float\[Array, 'a b'\], y: "
|
|
r"Float\[Array, 'b c'\]\)."
|
|
),
|
|
"The current values for each jaxtyping axis annotation are as follows.",
|
|
"a=2",
|
|
"b=3",
|
|
]
|
|
for match in matches:
|
|
with pytest.raises(TypeCheckError, match=match):
|
|
g(x, y=y)
|
|
|
|
|
|
def test_return(typecheck):
|
|
@jaxtyped(typechecker=typecheck)
|
|
def f(x: PyTree[Any, " T"], y: PyTree[Any, " S"]) -> PyTree[Any, "T S"]:
|
|
return "foo"
|
|
|
|
x = (1, 2)
|
|
y = {"a": 1}
|
|
matches = [
|
|
"Type-check error whilst checking the return value of f",
|
|
r"Called with parameters: {'x': \(1, 2\), 'y': {'a': 1}}",
|
|
"Actual value: 'foo'",
|
|
r"Expected type: PyTree\[Any, \"T S\"\].",
|
|
(
|
|
"The current values for each jaxtyping PyTree structure annotation are as "
|
|
"follows."
|
|
),
|
|
r"T=PyTreeDef\(\(\*, \*\)\)",
|
|
r"S=PyTreeDef\({'a': \*}\)",
|
|
]
|
|
for match in matches:
|
|
with pytest.raises(TypeCheckError, match=match):
|
|
f(x, y=y)
|
|
|
|
|
|
def test_dataclass_attribute(typecheck):
|
|
@jaxtyped(typechecker=typecheck)
|
|
class M(eqx.Module):
|
|
x: Float[Array, " *foo"]
|
|
y: PyTree[Any, " T"]
|
|
z: int
|
|
|
|
x = jnp.zeros((2, 3))
|
|
y = (1, (3, 4))
|
|
z = "not-an-int"
|
|
|
|
matches = [
|
|
"Type-check error whilst checking the parameters of M",
|
|
"The problem arose whilst typechecking parameter 'z'.",
|
|
(
|
|
r"Called with parameters: {'self': M\(\.\.\.\), 'x': f32\[2,3\], "
|
|
r"'y': \(1, \(3, 4\)\), 'z': 'not-an-int'}"
|
|
),
|
|
(
|
|
r"Parameter annotations: \(self: Any, x: Float\[Array, '\*foo'\], "
|
|
r"y: PyTree\[Any, \"T\"\], z: int\)."
|
|
),
|
|
"The current values for each jaxtyping axis annotation are as follows.",
|
|
r"foo=\(2, 3\)",
|
|
(
|
|
"The current values for each jaxtyping PyTree structure annotation are as "
|
|
"follows."
|
|
),
|
|
r"T=PyTreeDef\(\(\*, \(\*, \*\)\)\)",
|
|
]
|
|
for match in matches:
|
|
with pytest.raises(TypeCheckError, match=match):
|
|
M(x, y, z)
|