diff --git a/API.md b/API.md index d82e65d..e043b50 100644 --- a/API.md +++ b/API.md @@ -2,67 +2,83 @@ ## Annotating array types -Each array is denoted by a type `dtype[shape]`, such as `f32["batch channels"]`. +Each array is denoted by a type `dtype[array, shape]`, such as `Float[jnp.ndarray, "batch channels"]`. ### Shape The shape should be a string of space-separated symbols, such as "a b c d". Each symbol can be either an: -- `int`: fixed-size axis, e.g. `f32["28 28"]`. -- `str`: variable-size axis, e.g. `f32["channels"]`. -- A symbolic expression (without spaces!) in terms of other variable-size axes, e.g. `def remove_last(x: f32["dim"]) -> f32["dim-1"]`. +- `int`: fixed-size axis, e.g. `"28 28"`. +- `str`: variable-size axis, e.g. `"channels"`. +- A symbolic expression (without spaces!) in terms of other variable-size axes, e.g. `def remove_last(x: Float[jnp.ndarray, "dim"]) -> Float[jnp.ndarray, "dim-1"]`. When calling a function, variable-size axes and symbolic axes will be matched up across all arguments and checked for consistency. (See [runtime type checking](#runtime-type-checking) below.) In addition some modifiers can be applied: -- Prepend `*` to a dimension to indicate that it can match multiple axes, e.g. `f32["*batch c h w"]` will match zero or more batch axes. -- Prepend `#` to a dimension to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable, e.g. `add(x: f32["#foo"], y: f32["#foo"]) -> f32["#foo"]`. -- Prepend `_` to a dimension to disable any runtime checking of that dimension (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `f32["b c _ _"]`. +- Prepend `*` to a dimension to indicate that it can match multiple axes, e.g. `"*batch c h w"` will match zero or more batch axes. +- Prepend `#` to a dimension to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable, e.g. `add(x: Float[jnp.ndarray, "#foo"], y: Float[jnp.ndarray, "#foo"]) -> Float[jnp.ndarray, "#foo"]`. +- Prepend `_` to a dimension to disable any runtime checking of that dimension (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `"b c _ _"`. The order of these modifiers does not matter. As a special case: -- `...`: anonymous zero or more axes (equivalent to `*_`) e.g. `f32["... c h w"]` +- `...`: anonymous zero or more axes (equivalent to `*_`) e.g. `"... c h w"` Some notes: -- To denote a scalar shape use `""`, e.g. `f32[""]`. -- To denote an arbitrary shape (and only check dtype) use `"..."`, e.g. `f32["..."]`. +- To denote a scalar shape use `""`, e.g. `Float[jnp.ndarray, ""]`. +- To denote an arbitrary shape (and only check dtype) use `"..."`, e.g. `Float[jnp.ndarray, "..."]`. - You cannot have more than one use of multiple-axes, i.e. you can only use `...` or `*name` at most once in each array. -- An example of broadcasting multiple dimensions: `add(x: f32["*#foo"], y: f32["*#foo"]) -> f32["*#foo"]`. +- An example of broadcasting multiple dimensions: `add(x: Float[jnp.ndarray, "*#foo"], y: Float[jnp.ndarray, "*#foo"]) -> Float[jnp.ndarray, "*#foo"]`. - A symbolic expression cannot be evaluated unless all of the axes sizes it refers to have already been processed. In practice this usually means that they should only be used in annotations for the return type, and only use axes declared in the arguments. ### Dtype The dtype should be any one of (imported from `jaxtyping`): -- Any dtype at all: `Array` - - Boolean: `b` - - Any integer, unsigned integer, floating, or complex: `n` (for number) - - Any floating or complex: `x` (for inexact) - - Any floating point: `f` - - Floating point: `bf16`, `f16`, `f32`, `f64` (`bf16` is bfloat16) - - Any complex: `c` - - Complexes: `c64`, `c128` - - Any integer or unsigned intger: `t` (for integer) - - Any unsigned integer: `u` - - Unsigned integer: `u8`, `u16`, `u32`, `u64` - - Any signed integer: `i` - - Signed integer: `i8`, `i16`, `i32`, `i64` +- Any dtype at all: `Shaped` + - Boolean: `Bool` + - Any integer, unsigned integer, floating, or complex: `Num` + - Any floating or complex: `Inexact` + - Any floating point: `Float` + - Of particular precision: `bf16`, `f16`, `f32`, `f64` (`bf16` is bfloat16) + - Any complex: `Complex` + - Of particular precision: `c64`, `c128` + - Any integer or unsigned intger: `Int` + - Any unsigned integer: `IntUnsign` + - Of particular precision: `u8`, `u16`, `u32`, `u64` + - Any signed integer: `IntSign` + - Of particular precision: `i8`, `i16`, `i32`, `i64` Unless you really want to force a particular precision, then for most applications you should probably allow any floating-point, any integer, etc. That is, use ```python -from jaxtyping import f -f["some_shape"] +from jaxtyping import Float +Float[jnp.ndarray, "some_shape"] ``` rather than ```python from jaxtyping import f32 -f32["some_shape"] +f32[jnp.ndarray, "some_shape"] +``` + +### Array + +The array should typically be a `jnp.ndarray`. In practice, to save a bit of space and because it looks quite nice, we recommend: +```python +from jax.numpy import ndarray as Array +Float[Array, "..."] +``` + +You can use other types as well. `jaxtyping` has support for JAX, NumPy, TensorFlow, and PyTorch, e.g.: +```python +Float[jnp.ndarray, "..."] +Float[np.ndarray, "..."] +Float[tf.Tensor, "..."] +Float[torch.Tensor, "..."] ``` ## PyTrees ### `jaxtyping.PyTree` -Each PyTree is denoted by a type `PyTree[LeafType]`, such as `PyTree[int]` or `PyTree[Union[str, f32["b c"]]]`. +Each PyTree is denoted by a type `PyTree[LeafType]`, such as `PyTree[int]` or `PyTree[Union[str, f32[jnp.ndarray, "b c"]]]`. You can leave off the `[...]`, in which case `PyTree` is simply a suggestively-named alternative to `Any`. ([By definition all types are PyTrees.](https://jax.readthedocs.io/en/latest/pytrees.html)) @@ -83,6 +99,7 @@ Example: ```python # Import both the annotation and the `jaxtyped` decorator from `jaxtyping` from jaxtyping import f32, jaxtyped +from jax.numpy import ndarray as Array # Use your favourite typechecker: usually one of the two lines below. from typeguard import typechecked as typechecker @@ -91,7 +108,9 @@ from beartype import beartype as typechecker # Write your function. @jaxtyped must be applied above @typechecker! @jaxtyped @typechecker -def batch_outer_product(x: f32["b c1"], y: f32["b c2"]) -> f32["b c1 c2"]: +def batch_outer_product(x: f32[Array, "b c1"], + y: f32[Array, "b c2"] + ) -> f32[Array, "b c1 c2"]: return x[:, :, None] * y[:, None, :] ``` @@ -136,9 +155,14 @@ Any module imported **afterwards**, whose name begins with the specified string, The import hook may be uninstalled after you've imported all the modules you're interested in: ```python +# Manual uninstall hook = install_import_hook(...) ... # perform imports hook.uninstall() + +# Alternative: automatic uninstall +with install_import_hook(...): + # perform imports ``` The import hook can be applied to multiple packages via @@ -156,8 +180,9 @@ import do_stuff ### do_stuff.py from jaxtyping import f32 +from jax.numpy import ndarray as Array -def g(x: f32["..."]): +def g(x: f32[Array, "..."]): ... ``` @@ -166,11 +191,10 @@ def g(x: f32["..."]): ```python ### __init__.py from jaxtyping import install_import_hook -hook = install_import_hook("my_library_name", ("beartype", "beartype")) -from .subpackage import foo # full name is my_library_name.subpackage so will be hook'd -from .another_subpackage import bar # full name is my_library_name.another_subpackage so will be hook'd. -hook.uninstall() -del hook, install_import_hook, jaxtyping # keep interface tidy +with install_import_hook("my_library_name", ("beartype", "beartype")): + from .subpackage import foo # full name is my_library_name.subpackage so will be hook'd + from .another_subpackage import bar # full name is my_library_name.another_subpackage so will be hook'd. +del install_import_hook # keep interface tidy ``` #### pytest hook @@ -200,4 +224,4 @@ Union[u8["shape"], u16["shape"]] ### `jaxtyping.AbstractArray` The base class of all shape-and-dtype-specified arrays, e.g. it's a base class -for `f32["foo"]`. +for `f32[jnp.ndarray, "foo"]`. diff --git a/README.md b/README.md index 9bc0e93..ac6a0df 100644 --- a/README.md +++ b/README.md @@ -7,15 +7,18 @@ Type annotations **and runtime checking** for: **For example:** ```python -from jaxtyping import f32, PyTree +from jaxtyping import Float, PyTree +from jax.numpy import ndarray as Array -def matrix_multiply(x: f32["dim1 dim2"], y: f32["dim2 dim3"]) -> f32["dim1 dim3"]: +def matrix_multiply(x: Float[Array, "dim1 dim2"], + y: Float[Array, "dim2 dim3"] + ) -> Float[Array, "dim1 dim3"]: ... def accepts_pytree_of_ints(x: PyTree[int]): ... -def accepts_pytree_of_arrays(x: PyTree[f32["batch c1 c2"]]): +def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]): ... ``` @@ -49,7 +52,7 @@ SymPy<->JAX conversion; train symbolic expressions via gradient descent: [sympy2 Shape annotations + runtime type checking is inspired by [TorchTyping](https://github.com/patrick-kidger/torchtyping). -The concise syntax is inspired by [etils.array_types](https://github.com/google/etils/tree/main/etils/array_types). +The concise syntax is partially inspired by [etils.array_types](https://github.com/google/etils/tree/main/etils/array_types). ### Disclaimer diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 41a8428..dfa38d2 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -23,21 +23,30 @@ from .array_types import ( Array, b, bf16, + Bool, c, c64, c128, + Complex, f, f16, f32, f64, + Float, get_array_name_format, i, i8, i16, i32, i64, + Inexact, + Int, + IntSign, + IntUnsign, n, + Num, set_array_name_format, + Shaped, t, u, u8, @@ -51,4 +60,4 @@ from .import_hook import install_import_hook from .pytree_type import PyTree -__version__ = "0.1.0" +__version__ = "0.2.0" diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 8b62f22..9c134a0 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/array_types.py @@ -19,10 +19,9 @@ import enum import functools as ft -from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union +from typing import Any, Dict, List, NoReturn, Optional, Tuple, TYPE_CHECKING, Union from typing_extensions import Literal -import jax.numpy as jnp import numpy as np from .decorator import storage @@ -128,10 +127,26 @@ def _check_dims( class _MetaAbstractArray(type): def __instancecheck__(cls, obj): - if not isinstance(obj, jnp.ndarray): + if not isinstance(obj, cls.array_type): return False - if cls.dtypes is not _any_dtype and obj.dtype not in cls.dtypes: + if hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"): + # JAX, numpy + dtype = obj.dtype.type.__name__ + elif hasattr(obj.dtype, "as_numpy_dtype"): + # TensorFlow + dtype = obj.dtype.as_numpy_dtype.__name__ + else: + # PyTorch + repr_dtype = repr(obj.dtype).split(".") + if len(repr_dtype) == 2 and repr_dtype[0] == "torch": + dtype = repr_dtype[1] + else: + raise RuntimeError( + "Unrecognised array/tensor type to extract dtype from" + ) + + if cls.dtypes is not _any_dtype and dtype not in cls.dtypes: return False if len(storage.memo_stack) == 0: @@ -219,7 +234,8 @@ class _MetaAbstractArray(type): class AbstractArray(metaclass=_MetaAbstractArray): - dtypes: List[jnp.dtype] + array_type: Any + dtypes: List[str] dims: List[_AbstractDimOrVariadicDim] index_variadic: Optional[int] @@ -227,13 +243,25 @@ class AbstractArray(metaclass=_MetaAbstractArray): class _MetaAbstractDtype(type): def __instancecheck__(cls, obj: Any) -> NoReturn: raise RuntimeError( - f"Do not use `isinstance(x, jaxtyping.{cls.__name__}`. If you want to " + f"Do not use `isinstance(x, jaxtyping.{cls.__name__})`. If you want to " "check just the dtype of an array, then use " - f'`jaxtyping.{cls.__name__}["..."]`.' + f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.' ) @ft.lru_cache(maxsize=None) - def __getitem__(cls, dim_str: str) -> _MetaAbstractArray: + def __getitem__(cls, item: Tuple[Any, str]) -> _MetaAbstractArray: + if cls.deprecated is not None: + raise ValueError( + f"As of jaxtyping v0.2.0, {cls.__name__} has been deprecated in favour " + f"of {cls.deprecated.__name__}" + ) + if not isinstance(item, tuple) or len(item) != 2: + raise ValueError( + "As of jaxtyping v0.2.0, type annotations must now include an explicit " + "array type. For example `jaxtyping.f32[jnp.ndarray, 'foo bar']`." + ) + array_type, dim_str = item + del item if not isinstance(dim_str, str): raise ValueError( "Shape specification must be a string. Axes should be separated with " @@ -360,7 +388,7 @@ class _MetaAbstractDtype(type): elem = _SymbolicDim(elem, broadcastable) dims.append(elem) if _array_name_format == "dtype_and_shape": - name = f"{cls.__name__}['{dim_str}']" + name = f"{cls.__name__}[{array_type.__name__}, '{dim_str}']" elif _array_name_format == "array": name = "Array" else: @@ -368,88 +396,126 @@ class _MetaAbstractDtype(type): return _MetaAbstractArray( name, (AbstractArray,), - dict(dtypes=cls.dtypes, dims=dims, index_variadic=index_variadic), + dict( + array_type=array_type, + dtypes=cls.dtypes, + dims=dims, + index_variadic=index_variadic, + ), ) class AbstractDtype(metaclass=_MetaAbstractDtype): - dtypes: Union[str, List[str], Literal[_any_dtype]] + deprecated: Optional[str] + dtypes: Union[Literal[_any_dtype], List[str]] def __init__(self, *args, **kwargs): raise RuntimeError( "AbstractDtype cannot be instantiated. Perhaps you wrote e.g. " - '`f32("shape")` when you mean `f32["shape"]`?' + '`f32("shape")` when you mean `f32[jnp.ndarray, "shape"]`?' ) def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) - dtypes = cls.dtypes - if dtypes is not _any_dtype: - if not isinstance(dtypes, list): - dtypes = [dtypes] - dtypes = [jnp.dtype(d) for d in dtypes] + dtypes: Union[Literal[_any_dtype], str, List[str]] = cls.dtypes + if isinstance(dtypes, str): + dtypes = [dtypes] cls.dtypes = dtypes -_bool = "bool" -_uint8 = "uint8" -_uint16 = "uint16" -_uint32 = "uint32" -_uint64 = "uint64" -_int8 = "int8" -_int16 = "int16" -_int32 = "int32" -_int64 = "int64" -_bfloat16 = "bfloat16" -_float16 = "float16" -_float32 = "float32" -_float64 = "float64" -_complex64 = "complex64" -_complex128 = "complex128" +if TYPE_CHECKING: + # Note that `from typing_extensions import Annotated; ... = Annotated` + # does not work with static type checkers. `Annotated` is a typeform rather + # than a type, meaning it cannot be assigned. + from typing_extensions import Annotated as Bool + from typing_extensions import Annotated as Complex + from typing_extensions import Annotated as Float + from typing_extensions import Annotated as Inexact + from typing_extensions import Annotated as Int + from typing_extensions import Annotated as IntSign + from typing_extensions import Annotated as IntUnsign + from typing_extensions import Annotated as Num + from typing_extensions import Annotated as Shaped + from typing_extensions import Annotated as bf16 + from typing_extensions import Annotated as c64 + from typing_extensions import Annotated as c128 + from typing_extensions import Annotated as f16 + from typing_extensions import Annotated as f32 + from typing_extensions import Annotated as f64 + from typing_extensions import Annotated as i8 + from typing_extensions import Annotated as i16 + from typing_extensions import Annotated as i32 + from typing_extensions import Annotated as i64 + from typing_extensions import Annotated as u8 + from typing_extensions import Annotated as u16 + from typing_extensions import Annotated as u32 + from typing_extensions import Annotated as u64 +else: + _bool = "bool" + _uint8 = "uint8" + _uint16 = "uint16" + _uint32 = "uint32" + _uint64 = "uint64" + _int8 = "int8" + _int16 = "int16" + _int32 = "int32" + _int64 = "int64" + _bfloat16 = "bfloat16" + _float16 = "float16" + _float32 = "float32" + _float64 = "float64" + _complex64 = "complex64" + _complex128 = "complex128" + def _make_dtype(_dtypes, name, *, _deprecated=None): + class _Cls(AbstractDtype): + deprecated = _deprecated + dtypes = _dtypes -def _make_dtype(_dtypes, name): - class _Cls(AbstractDtype): - dtypes = _dtypes + _Cls.__name__ = name + _Cls.__qualname__ = name + return _Cls - _Cls.__name__ = name - _Cls.__qualname__ = name - return _Cls + Bool = _make_dtype(_bool, "Bool") + u8 = _make_dtype(_uint8, "u8") + u16 = _make_dtype(_uint16, "u16") + u32 = _make_dtype(_uint32, "u32") + u64 = _make_dtype(_uint64, "u64") + i8 = _make_dtype(_int8, "i8") + i16 = _make_dtype(_int16, "i16") + i32 = _make_dtype(_int32, "i32") + i64 = _make_dtype(_int64, "i64") + bf16 = _make_dtype(_bfloat16, "bf16") + f16 = _make_dtype(_float16, "f16") + f32 = _make_dtype(_float32, "f32") + f64 = _make_dtype(_float64, "f64") + c64 = _make_dtype(_complex64, "c64") + c128 = _make_dtype(_complex128, "c128") + uints = [_uint8, _uint16, _uint32, _uint64] + ints = [_int8, _int16, _int32, _int64] + floats = [_bfloat16, _float16, _float32, _float64] + complexes = [_complex64, _complex128] -b = _make_dtype(_bool, "b") -u8 = _make_dtype(_uint8, "u8") -u16 = _make_dtype(_uint16, "u16") -u32 = _make_dtype(_uint32, "u32") -u64 = _make_dtype(_uint64, "u64") -i8 = _make_dtype(_int8, "i8") -i16 = _make_dtype(_int16, "i16") -i32 = _make_dtype(_int32, "i32") -i64 = _make_dtype(_int64, "i64") -bf16 = _make_dtype(_bfloat16, "bf16") -f16 = _make_dtype(_float16, "f16") -f32 = _make_dtype(_float32, "f32") -f64 = _make_dtype(_float64, "f64") -c64 = _make_dtype(_complex64, "c64") -c128 = _make_dtype(_complex128, "c128") + # We match NumPy's type hierarachy in what types to provide. See the diagram at + # https://numpy.org/doc/stable/reference/arrays.scalars.html#scalars -uints = [_uint8, _uint16, _uint32, _uint64] -ints = [_int8, _int16, _int32, _int64] -floats = [_bfloat16, _float16, _float32, _float64] -complexes = [_complex64, _complex128] + IntUnsign = _make_dtype(uints, "IntUnsign") + IntSign = _make_dtype(ints, "IntSign") + Int = _make_dtype(uints + ints, "Int") + Float = _make_dtype(floats, "Float") + Complex = _make_dtype(complexes, "Complex") + Inexact = _make_dtype(floats + complexes, "Inexact") # inexact + Num = _make_dtype(uints + ints + floats + complexes, "Num") # number + Shaped = _make_dtype(_any_dtype, "Shaped") -# We match NumPy's type hierarachy in what types to provide. See the diagram at -# https://numpy.org/doc/stable/reference/arrays.scalars.html#scalars -# -# No attempt is made to match up against their character codes: all of the below are -# abstract base classes without NumPy chararacter codes. - -u = _make_dtype(uints, "u") -i = _make_dtype(ints, "i") -t = _make_dtype(uints + ints, "t") # integer -f = _make_dtype(floats, "f") -c = _make_dtype(complexes, "c") -x = _make_dtype(floats + complexes, "x") # inexact -n = _make_dtype(uints + ints + floats + complexes, "n") # number -Array = _make_dtype(_any_dtype, "Array") + b = _make_dtype(_bool, "b", _deprecated=Bool) + i = _make_dtype(_bool, "i", _deprecated=IntSign) + u = _make_dtype(_bool, "u", _deprecated=IntUnsign) + t = _make_dtype(_bool, "t", _deprecated=Int) + f = _make_dtype(_bool, "f", _deprecated=Float) + c = _make_dtype(_bool, "c", _deprecated=Complex) + x = _make_dtype(_bool, "x", _deprecated=Inexact) + n = _make_dtype(_bool, "n", _deprecated=Num) + Array = _make_dtype(_bool, "Array", _deprecated=Shaped) diff --git a/test/import_hook_tester_beartype.py b/test/import_hook_tester_beartype.py index 6e69d53..11da8dd 100644 --- a/test/import_hook_tester_beartype.py +++ b/test/import_hook_tester_beartype.py @@ -25,7 +25,7 @@ from jaxtyping import f32 from .helpers import ParamError -def g(x: f32[" b"]): +def g(x: f32[jnp.ndarray, " b"]): pass diff --git a/test/import_hook_tester_broken_checker.py b/test/import_hook_tester_broken_checker.py index 6e69d53..11da8dd 100644 --- a/test/import_hook_tester_broken_checker.py +++ b/test/import_hook_tester_broken_checker.py @@ -25,7 +25,7 @@ from jaxtyping import f32 from .helpers import ParamError -def g(x: f32[" b"]): +def g(x: f32[jnp.ndarray, " b"]): pass diff --git a/test/import_hook_tester_transitive/another_file.py b/test/import_hook_tester_transitive/another_file.py index b7c9a9b..9da8e5a 100644 --- a/test/import_hook_tester_transitive/another_file.py +++ b/test/import_hook_tester_transitive/another_file.py @@ -25,7 +25,7 @@ from jaxtyping import f32 from ..helpers import ParamError -def g(x: f32[" b"]): +def g(x: f32[jnp.ndarray, " b"]): pass diff --git a/test/import_hook_tester_typeguard.py b/test/import_hook_tester_typeguard.py index 6e69d53..11da8dd 100644 --- a/test/import_hook_tester_typeguard.py +++ b/test/import_hook_tester_typeguard.py @@ -25,7 +25,7 @@ from jaxtyping import f32 from .helpers import ParamError -def g(x: f32[" b"]): +def g(x: f32[jnp.ndarray, " b"]): pass diff --git a/test/test_array.py b/test/test_array.py index 7f94f24..c30b65e 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -21,15 +21,18 @@ import jax.numpy as jnp import jax.random as jr import pytest -from jaxtyping import Array, f, f32, jaxtyped +from jaxtyping import f32, Float, jaxtyped, Shaped from .helpers import ParamError, ReturnError +Array = jnp.ndarray + + def test_basic(typecheck): @jaxtyped @typecheck - def g(x: Array["..."]): + def g(x: Shaped[Array, "..."]): pass g(jnp.array(1.0)) @@ -38,14 +41,14 @@ def test_basic(typecheck): def test_return(typecheck, getkey): @jaxtyped @typecheck - def g(x: f["b c"]) -> f["c b"]: + def g(x: Float[Array, "b c"]) -> Float[Array, "c b"]: return jnp.transpose(x) g(jr.normal(getkey(), (3, 4))) @jaxtyped @typecheck - def h(x: f["b c"]) -> f["b c"]: + def h(x: Float[Array, "b c"]) -> Float[Array, "b c"]: return jnp.transpose(x) with pytest.raises(ReturnError): @@ -55,7 +58,7 @@ def test_return(typecheck, getkey): def test_two_args(typecheck, getkey): @jaxtyped @typecheck - def g(x: Array["b c"], y: Array["c d"]): + 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))) @@ -64,7 +67,7 @@ def test_two_args(typecheck, getkey): @jaxtyped @typecheck - def h(x: Array["b c"], y: Array["c d"]) -> Array["b d"]: + 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))) @@ -75,7 +78,7 @@ def test_two_args(typecheck, getkey): def test_any_dtype(typecheck, getkey): @jaxtyped @typecheck - def g(x: Array["a b"]) -> Array["a b"]: + def g(x: Shaped[Array, "a b"]) -> Shaped[Array, "a b"]: return x g(jr.normal(getkey(), (3, 4))) @@ -92,12 +95,12 @@ def test_any_dtype(typecheck, getkey): def test_nested_jaxtyped(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32["b c"], transpose: bool) -> f32["c b"]: + def g(x: f32[Array, "b c"], transpose: bool) -> f32[Array, "c b"]: return h(x, transpose) @jaxtyped @typecheck - def h(x: f32["c b"], transpose: bool) -> f32["b c"]: + def h(x: f32[Array, "c b"], transpose: bool) -> f32[Array, "b c"]: if transpose: return jnp.transpose(x) else: @@ -113,11 +116,11 @@ def test_nested_jaxtyped(typecheck, getkey): def test_nested_nojaxtyped(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32["b c"]): + def g(x: f32[Array, "b c"]): return h(x) @typecheck - def h(x: f32["c b"]): + def h(x: f32[Array, "c b"]): return x with pytest.raises(ParamError): @@ -127,14 +130,14 @@ def test_nested_nojaxtyped(typecheck, getkey): def test_isinstance(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32["b c"]) -> f32[" z"]: + def g(x: f32[Array, "b c"]) -> f32[Array, " z"]: y = jnp.transpose(x) - assert isinstance(y, f32["c b"]) + assert isinstance(y, f32[Array, "c b"]) assert not isinstance( - y, f32["b z"] + y, f32[Array, "b z"] ) # z left unbound as b!=c (unless x symmetric, which it isn't) out = jr.normal(getkey(), (500,)) - assert isinstance(out, f32["z"]) # z now bound + assert isinstance(out, f32[Array, "z"]) # z now bound return out g(jr.normal(getkey(), (2, 3))) @@ -143,7 +146,7 @@ def test_isinstance(typecheck, getkey): def test_fixed(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32["4 5 foo"], y: f32[" foo"]) -> f32["4 5"]: + def g(x: f32[Array, "4 5 foo"], y: f32[Array, " foo"]) -> f32[Array, "4 5"]: return x @ y a = jr.normal(getkey(), (4, 5, 2)) @@ -158,7 +161,7 @@ def test_fixed(typecheck, getkey): def test_anonymous(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32["foo _"], y: f32[" _"]): + def g(x: f32[Array, "foo _"], y: f32[Array, " _"]): pass a = jr.normal(getkey(), (3, 4)) @@ -169,7 +172,7 @@ def test_anonymous(typecheck, getkey): def test_named_variadic(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32["*batch foo"], y: f32[" *batch"], z: f32[" foo"]): + def g(x: f32[Array, "*batch foo"], y: f32[Array, " *batch"], z: f32[Array, " foo"]): pass c = jr.normal(getkey(), (5,)) @@ -189,7 +192,7 @@ def test_named_variadic(typecheck, getkey): @jaxtyped @typecheck - def h(x: f32[" foo *batch"], y: f32[" foo *batch bar"]): + def h(x: f32[Array, " foo *batch"], y: f32[Array, " foo *batch bar"]): pass a = jr.normal(getkey(), (4,)) @@ -205,7 +208,7 @@ def test_named_variadic(typecheck, getkey): def test_anonymous_variadic(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32["... foo"], y: f32[" foo"]): + def g(x: f32[Array, "... foo"], y: f32[Array, " foo"]): pass a1 = jr.normal(getkey(), (5,)) @@ -227,7 +230,7 @@ def test_anonymous_variadic(typecheck, getkey): def test_broadcast_fixed(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32["#4"]): + def g(x: f32[Array, "#4"]): pass g(jr.normal(getkey(), (4,))) @@ -240,7 +243,7 @@ def test_broadcast_fixed(typecheck, getkey): def test_broadcast_named(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[" #foo"], y: f32[" #foo"]): + def g(x: f32[Array, " #foo"], y: f32[Array, " #foo"]): pass a = jr.normal(getkey(), (3,)) @@ -264,7 +267,7 @@ def test_broadcast_named(typecheck, getkey): def test_broadcast_variadic_named(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[" *#foo"], y: f32[" *#foo"]): + def g(x: f32[Array, " *#foo"], y: f32[Array, " *#foo"]): pass a = jr.normal(getkey(), (3,)) @@ -322,28 +325,28 @@ def test_broadcast_variadic_named(typecheck, getkey): def test_no_commas(): with pytest.raises(ValueError): - f32["foo, bar"] + f32[Array, "foo, bar"] def test_symbolic(typecheck, getkey): @jaxtyped @typecheck - def make_slice(x: f32[" dim"]) -> f32[" dim-1"]: + def make_slice(x: f32[Array, " dim"]) -> f32[Array, " dim-1"]: return x[1:] @jaxtyped @typecheck - def cat(x: f32[" dim"]) -> f32[" 2*dim"]: + def cat(x: f32[Array, " dim"]) -> f32[Array, " 2*dim"]: return jnp.concatenate([x, x]) @jaxtyped @typecheck - def bad_make_slice(x: f32[" dim"]) -> f32[" dim-1"]: + def bad_make_slice(x: f32[Array, " dim"]) -> f32[Array, " dim-1"]: return x @jaxtyped @typecheck - def bad_cat(x: f32[" dim"]) -> f32[" 2*dim"]: + def bad_cat(x: f32[Array, " dim"]) -> f32[Array, " 2*dim"]: return jnp.concatenate([x, x, x]) x = jr.normal(getkey(), (5,)) @@ -365,7 +368,7 @@ def test_symbolic(typecheck, getkey): def test_incomplete_symbolic(typecheck, getkey): @jaxtyped @typecheck - def foo(x: f32[" 2*dim"]): + def foo(x: f32[Array, " 2*dim"]): pass x = jr.normal(getkey(), (4,)) diff --git a/test/test_pytree.py b/test/test_pytree.py index 2bcbd7c..244ebaf 100644 --- a/test/test_pytree.py +++ b/test/test_pytree.py @@ -25,7 +25,7 @@ import jax.numpy as jnp import jax.random as jr import pytest -from jaxtyping import f, jaxtyped, PyTree +from jaxtyping import Float, jaxtyped, PyTree from .helpers import make_mlp, ParamError @@ -95,7 +95,7 @@ def test_nested_pytrees(getkey, typecheck): def test_pytree_array(typecheck): @jaxtyped @typecheck - def g(x: PyTree[f["..."]]): + def g(x: PyTree[Float[jnp.ndarray, "..."]]): pass g(jnp.array(1.0)) @@ -109,7 +109,7 @@ def test_pytree_array(typecheck): def test_pytree_shaped_array(typecheck, getkey): @jaxtyped @typecheck - def g(x: PyTree[f["b c"]]): + def g(x: PyTree[Float[jnp.ndarray, "b c"]]): pass g(jnp.array([[1.0]]))