mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-08-21 11:16:29 +08:00
152 lines
3.8 KiB
Python
152 lines
3.8 KiB
Python
import pytest
|
|
from IPython.testing.globalipapp import start_ipython
|
|
|
|
from .helpers import ParamError
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def session_ip():
|
|
yield start_ipython()
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def ip(session_ip):
|
|
session_ip.run_cell(raw_cell="import jaxtyping")
|
|
session_ip.run_line_magic(magic_name="load_ext", line="jaxtyping")
|
|
session_ip.run_line_magic(
|
|
magic_name="jaxtyping.typechecker", line="typeguard.typechecked"
|
|
)
|
|
yield session_ip
|
|
|
|
|
|
def test_that_ipython_works(ip):
|
|
ip.run_cell(raw_cell="x = 1").raise_error()
|
|
assert ip.user_global_ns["x"] == 1
|
|
|
|
|
|
def test_function_beartype(ip):
|
|
ip.run_cell(
|
|
raw_cell="""
|
|
def f(x: int):
|
|
pass
|
|
"""
|
|
).raise_error()
|
|
ip.run_cell(raw_cell="f(1)").raise_error()
|
|
|
|
with pytest.raises(ParamError):
|
|
ip.run_cell(raw_cell='f("x")').raise_error()
|
|
|
|
|
|
def test_function_none(ip):
|
|
ip.run_cell(
|
|
raw_cell="""
|
|
def f(a,b,c):
|
|
pass
|
|
"""
|
|
).raise_error()
|
|
ip.run_cell(raw_cell='f(1,2,"k")').raise_error()
|
|
|
|
|
|
def test_function_jaxtyped(ip):
|
|
ip.run_cell(
|
|
raw_cell="""
|
|
from jaxtyping import Float, Array, Int
|
|
import jax
|
|
|
|
def g(x: Float[Array, "1"]):
|
|
return x + 1
|
|
|
|
"""
|
|
).raise_error()
|
|
|
|
ip.run_cell(raw_cell="g(jax.numpy.array([1.0]))").raise_error()
|
|
|
|
with pytest.raises(ParamError):
|
|
ip.run_cell(raw_cell="g(jax.numpy.array(1.0))").raise_error()
|
|
|
|
with pytest.raises(ParamError):
|
|
ip.run_cell(raw_cell="g(jax.numpy.array([1]))").raise_error()
|
|
|
|
with pytest.raises(ParamError):
|
|
ip.run_cell(raw_cell="g(jax.numpy.array([2, 3]))").raise_error()
|
|
|
|
with pytest.raises(ParamError):
|
|
ip.run_cell(raw_cell='g("string")').raise_error()
|
|
|
|
|
|
def test_function_jaxtyped_and_jitted(ip):
|
|
ip.run_cell(
|
|
raw_cell="""
|
|
from jaxtyping import Float, Array, Int
|
|
import jax
|
|
|
|
@jax.jit
|
|
def g(x: Float[Array, "1"]):
|
|
return x + 1
|
|
|
|
"""
|
|
).raise_error()
|
|
|
|
ip.run_cell(raw_cell="g(jax.numpy.array([1.0]))").raise_error()
|
|
|
|
with pytest.raises(ParamError):
|
|
ip.run_cell(raw_cell="g(jax.numpy.array(1.0))").raise_error()
|
|
|
|
with pytest.raises(ParamError):
|
|
ip.run_cell(raw_cell="g(jax.numpy.array([1]))").raise_error()
|
|
|
|
with pytest.raises(ParamError):
|
|
ip.run_cell(raw_cell="g(jax.numpy.array([2, 3]))").raise_error()
|
|
|
|
with pytest.raises(ParamError):
|
|
ip.run_cell(raw_cell='g("string")').raise_error()
|
|
|
|
|
|
def test_class_jaxtyped(ip):
|
|
ip.run_cell(
|
|
raw_cell="""
|
|
from jaxtyping import Float, Array, Int
|
|
import equinox as eqx
|
|
import jax
|
|
|
|
class A(eqx.Module):
|
|
x: Float[Array, "2"]
|
|
|
|
def do_something(self, y: Int[Array, ""]):
|
|
return self.x + y
|
|
"""
|
|
).raise_error()
|
|
|
|
ip.run_cell(raw_cell="a = A(jax.numpy.array([1.0, 2.0]))").raise_error()
|
|
ip.run_cell(raw_cell="a.do_something(jax.numpy.array(2))").raise_error()
|
|
|
|
with pytest.raises(ParamError):
|
|
ip.run_cell(raw_cell="A(jax.numpy.array([1.0]))").raise_error()
|
|
|
|
with pytest.raises(ParamError):
|
|
ip.run_cell(
|
|
raw_cell="a.do_something(jax.numpy.array([2.0, 3.0]))"
|
|
).raise_error()
|
|
|
|
|
|
def test_class_not_dataclass(ip):
|
|
ip.run_cell(
|
|
raw_cell="""
|
|
from jaxtyping import Float, Array, Int
|
|
import equinox as eqx
|
|
import jax
|
|
|
|
class A:
|
|
def __init__(self, x):
|
|
self.x = x
|
|
|
|
def do_something(self, y):
|
|
return x + y
|
|
"""
|
|
).raise_error()
|
|
|
|
ip.run_cell(raw_cell="a = A(jax.numpy.array([1.0, 2.0]))").raise_error()
|
|
ip.run_cell(raw_cell="a.do_something(jax.numpy.array(2))").raise_error()
|
|
ip.run_cell(raw_cell="A(jax.numpy.array([1.0]))").raise_error()
|
|
ip.run_cell(raw_cell="a.do_something(jax.numpy.array([2.0, 3.0]))").raise_error()
|