From 903000f3d5bc2efe52a9be6b3d46572102f8f5be Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Tue, 30 Aug 2022 12:47:45 -0700 Subject: [PATCH 1/9] Rewrote syntax --- API.md | 96 +++++--- README.md | 11 +- jaxtyping/__init__.py | 11 +- jaxtyping/array_types.py | 208 ++++++++++++------ test/import_hook_tester_beartype.py | 2 +- test/import_hook_tester_broken_checker.py | 2 +- .../another_file.py | 2 +- test/import_hook_tester_typeguard.py | 2 +- test/test_array.py | 61 ++--- test/test_pytree.py | 6 +- 10 files changed, 253 insertions(+), 148 deletions(-) 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]])) From 48e81312470f4b2fc36b93d8e6dab72184e24c21 Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Tue, 30 Aug 2022 12:56:47 -0700 Subject: [PATCH 2/9] Added jaxtyping.Array=jnp.ndarray --- API.md | 39 ++++++++++++++++----------------------- README.md | 3 +-- jaxtyping/array_types.py | 9 +++++++++ test/test_array.py | 5 +---- 4 files changed, 27 insertions(+), 29 deletions(-) diff --git a/API.md b/API.md index e043b50..3bf8b4a 100644 --- a/API.md +++ b/API.md @@ -2,20 +2,20 @@ ## Annotating array types -Each array is denoted by a type `dtype[array, shape]`, such as `Float[jnp.ndarray, "batch channels"]`. +Each array is denoted by a type `dtype[array, shape]`, such as `Float[Array, "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. `"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"]`. +- A symbolic expression (without spaces!) in terms of other variable-size axes, e.g. `def remove_last(x: Float[Array, "dim"]) -> Float[Array, "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. `"*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 indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable, e.g. `add(x: Float[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#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. @@ -24,10 +24,10 @@ As a special case: - `...`: anonymous zero or more axes (equivalent to `*_`) e.g. `"... c h w"` Some notes: -- 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, "..."]`. +- To denote a scalar shape use `""`, e.g. `Float[Array, ""]`. +- To denote an arbitrary shape (and only check dtype) use `"..."`, e.g. `Float[Array, "..."]`. - 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: Float[jnp.ndarray, "*#foo"], y: Float[jnp.ndarray, "*#foo"]) -> Float[jnp.ndarray, "*#foo"]`. +- An example of broadcasting multiple dimensions: `add(x: Float[Array, "*#foo"], y: Float[Array, "*#foo"]) -> Float[Array, "*#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 @@ -49,26 +49,21 @@ The dtype should be any one of (imported from `jaxtyping`): 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 Float -Float[jnp.ndarray, "some_shape"] +from jaxtyping import Array, Float +Float[Array, "some_shape"] ``` rather than ```python -from jaxtyping import f32 -f32[jnp.ndarray, "some_shape"] +from jaxtyping import Array, f32 +f32[Array, "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, "..."] -``` +The array should typically be a `jaxtyping.Array`, which is an alias for `jnp.ndarray`. -You can use other types as well. `jaxtyping` has support for JAX, NumPy, TensorFlow, and PyTorch, e.g.: +But 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, "..."] @@ -78,7 +73,7 @@ Float[torch.Tensor, "..."] ### `jaxtyping.PyTree` -Each PyTree is denoted by a type `PyTree[LeafType]`, such as `PyTree[int]` or `PyTree[Union[str, f32[jnp.ndarray, "b c"]]]`. +Each PyTree is denoted by a type `PyTree[LeafType]`, such as `PyTree[int]` or `PyTree[Union[str, f32[Array, "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)) @@ -98,8 +93,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 +from jaxtyping import Array, f32, jaxtyped # Use your favourite typechecker: usually one of the two lines below. from typeguard import typechecked as typechecker @@ -179,8 +173,7 @@ install_import_hook("do_stuff", ("typeguard", "typechecked")) import do_stuff ### do_stuff.py -from jaxtyping import f32 -from jax.numpy import ndarray as Array +from jaxtyping import Array, f32 def g(x: f32[Array, "..."]): ... @@ -224,4 +217,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[jnp.ndarray, "foo"]`. +for `f32[Array, "foo"]`. diff --git a/README.md b/README.md index ac6a0df..df14556 100644 --- a/README.md +++ b/README.md @@ -7,8 +7,7 @@ Type annotations **and runtime checking** for: **For example:** ```python -from jaxtyping import Float, PyTree -from jax.numpy import ndarray as Array +from jaxtyping import Array, Float, PyTree def matrix_multiply(x: Float[Array, "dim1 dim2"], y: Float[Array, "dim2 dim3"] diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 9c134a0..7afbb37 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/array_types.py @@ -22,6 +22,7 @@ import functools as ft 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 @@ -261,6 +262,8 @@ class _MetaAbstractDtype(type): "array type. For example `jaxtyping.f32[jnp.ndarray, 'foo bar']`." ) array_type, dim_str = item + if array_type is Array: + array_type = jnp.ndarray del item if not isinstance(dim_str, str): raise ValueError( @@ -451,6 +454,8 @@ if TYPE_CHECKING: from typing_extensions import Annotated as u16 from typing_extensions import Annotated as u32 from typing_extensions import Annotated as u64 + + from jax.numpy import ndarray as Array else: _bool = "bool" _uint8 = "uint8" @@ -518,4 +523,8 @@ else: c = _make_dtype(_bool, "c", _deprecated=Complex) x = _make_dtype(_bool, "x", _deprecated=Inexact) n = _make_dtype(_bool, "n", _deprecated=Num) + # Note that Array also has a non-deprecated use-case as + # `f32[Array, "foo"]`. + # TODO: once these deprecations are removed, then just have a + # `from jax.numpy import ndarray as Array` in `__init__.py`. Array = _make_dtype(_bool, "Array", _deprecated=Shaped) diff --git a/test/test_array.py b/test/test_array.py index c30b65e..bcdc4c0 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -21,14 +21,11 @@ import jax.numpy as jnp import jax.random as jr import pytest -from jaxtyping import f32, Float, jaxtyped, Shaped +from jaxtyping import Array, f32, Float, jaxtyped, Shaped from .helpers import ParamError, ReturnError -Array = jnp.ndarray - - def test_basic(typecheck): @jaxtyped @typecheck From 07e735797e89b4424130f5a4d625a21e63b4178e Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Tue, 30 Aug 2022 13:14:50 -0700 Subject: [PATCH 3/9] Doc updates --- API.md | 10 +++++----- FAQ.md | 42 +++++++++++++++++++----------------------- README.md | 1 + 3 files changed, 25 insertions(+), 28 deletions(-) diff --git a/API.md b/API.md index 3bf8b4a..75f976a 100644 --- a/API.md +++ b/API.md @@ -127,7 +127,7 @@ and just do your own manual `isinstance` checks.) Only `isinstance` checks that pass will contribute to the store of axis name-size pairs; those that fail will not. As such it is safe to write e.g. `assert not isinstance(x, -f32["foo"])`. +f32[Array, "foo"])`. ### `jaxtyping.install_import_hook` @@ -156,7 +156,7 @@ hook.uninstall() # Alternative: automatic uninstall with install_import_hook(...): - # perform imports + ... # perform imports ``` The import hook can be applied to multiple packages via @@ -202,16 +202,16 @@ which will apply the import hook to all modules whose names start with either `f ### `jaxtyping.AbstractDtype` -The base class of all dtypes. This can be used to create your own custom collection of dtypes (analogous to `n`, `x` etc.) For example: +The base class of all dtypes. This can be used to create your own custom collection of dtypes (analogous to `Float`, `Inexact` etc.) For example: ```python class u8_or_u16(AbstractDtype): dtypes = ["uint8", "uint16"] -u8_or_u16["shape"] +u8_or_u16[Array, "shape"] ``` which is functionally equivalent to ```python -Union[u8["shape"], u16["shape"]] +Union[u8[Array, "shape"], u16[Array, "shape"]] ``` ### `jaxtyping.AbstractArray` diff --git a/FAQ.md b/FAQ.md index 1c73509..b62f847 100644 --- a/FAQ.md +++ b/FAQ.md @@ -6,38 +6,34 @@ In type annotations, strings are used for two different things. Sometimes they'r Some tooling in the Python ecosystem assumes that only the latter is true, and will throw spurious errors if you try to use a string just as a string (like we do). -In the case of `flake8`, at least, this is easily resolved. Multi-dimensional arrays (e.g. `f32["b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `f32["x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `f32[" x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before. +In the case of `flake8`, at least, this is easily resolved. Multi-dimensional arrays (e.g. `f32[Array, "b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `f32[Array, "x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `f32[Array, " x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before. -## What about support for static type checkers, like `mypy`, `pyright`, etc.? +## Does jaxtyping work with static type checkers like `mypy`/`pyright`/`pytype`? -Nope. +There is partial support for these. An annotation of the form `dtype[array, shape]` should be treated as just `array` by a static type checker. Unfortunately full dtype/shape checking is beyond the scope of what static type checking is currently capable of. -Python's static typing ecosystem is a complicated collection of edge cases. Many of them block ML/scientific computing in particular. A few examples: +(Note that at time of writing, `pytype` has a bug in that `dtype[array, shape]` is sometimes treated as `Any` rather than `array`. The other two work fine.) + +## How does jaxtyping interact with `jax.jit`? + +jaxtyping and `jax.jit` synergise beautifully. + +When calling JAX operations wrapped in a `jax.jit`, then the dtype/shape-checking will happen at trace time. (When JAX traces your function prior to compiling it.) The actual compiled code does not have any dtype/shape-checking, and will therefore still be just as fast as before! + +## Does jaxtyping use [PEP 646](https://www.python.org/dev/peps/pep-0646/) (variadic generics)? + +The intention of PEP 646 was to make it possible for static type checkers to perform shape checks of arrays. Unfortunately, this still isn't yet practical, so jaxtyping deliberately does not use this. (Yet?) + +The real problem is that Python's static typing ecosystem is a complicated collection of edge cases. Many of them block ML/scientific computing in particular. For example: 1. The static type system is intrinsically not expressive enough to describe operations like concatenation, stacking, or broadcasting. -2. Axes have to be lifted to type-level variables. Meanwhile the approach taken in libraries like `jaxtyping` and [TorchTyping](https://github.com/patrick-kidger/torchtyping) is to use value-level variables for types: because that's what the underlying JAX, PyTorch etc. libraries use! As such, making a static type checker work with these libraries would require either fundamentally rewriting these libraries, or exhaustively maintaining type stubs for them, and would *still* require a `typing.cast` any time you use anything unstubbed (e.g. any third party library, or part of your codebase you haven't typed yet). This is a huge maintenance burden for anyone. +2. Axes have to be lifted to type-level variables. Meanwhile the approach taken in libraries like `jaxtyping` and [TorchTyping](https://github.com/patrick-kidger/torchtyping) is to use value-level variables for types: because that's what the underlying JAX, PyTorch etc. libraries use! As such, making a static type checker work with these libraries would require either fundamentally rewriting these libraries, or exhaustively maintaining type stubs for them, and would *still* require a `typing.cast` any time you use anything unstubbed (e.g. any third party library, or part of your codebase you haven't typed yet). This is a huge maintenance burden. 3. Static type checkers have a variety of bugs that affect this use case. `mypy` doesn't support `Protocol`s correctly. `pyright` doesn't support genericised subprotocols. etc. -4. Variadic generics exist. Variadic protocols do not. (It's not clear that these have been contemplated.) +4. Variadic generics exist. Variadic protocols do not. (It's not clear that these were contemplated.) -5. The syntax for static typing is verbose. You have to write things like `Array[Unpack[AnyShape], Literal[3], Height, Width]` instead of `Array["... 3 height width"]`. +5. The syntax for static typing is verbose. You have to write things like `Array[Float32, Unpack[AnyShape], Literal[3], Height, Width]` instead of `f32[Array, "... 3 height width"]`. 6. [The underlying type system has flaws](https://github.com/patrick-kidger/torchtyping/issues/37#issuecomment-1153294196). [The numeric tower is broken](https://stackoverflow.com/a/69383462); [int is not a number](https://github.com/python/mypy/issues/3186#issuecomment-885718629); [virtual base classes don't work](https://github.com/python/mypy/issues/2922); [complex lies about having comparison operations, so type checkers have to lie about that lie in order to remove them again](https://posita.github.io/numerary/0.4/whytho/); `typing.*` don't work with `isinstance`; co/contra-variance are baked into containers (not specified at use-time); `dict` is variadic despite... not being variadic; bool is a subclass of int (!); ... etc. etc. - -## What about [PEP 646](https://www.python.org/dev/peps/pep-0646/) and variadic generics? - -[Doesn't change the previous issues, unfortunately.](https://github.com/patrick-kidger/torchtyping/issues/37) All the problems of the previous heading still hold true. They're just also true for types like `AnyDimensionalArray[Batch, Channels, AsManyArgumentsAsWePlease]` as well as types like `TwoDimensionalArray[Batch, Channels]`. - -## Is the lack of interaction with static typing a problem? - -At least for any software that is mostly just running JAX code, no! - -The correct way to use JAX is to put together all your operations, and then put a single `jax.jit` right at the very top. This gives you optimal speed; anything else will be unnecessarily (and substantially) slower. - -This means that all the type checking only gets resolved once: at trace time. Afterwards JAX still lowers everything down to the same optimised code. - -In some sense, `python myprogram.py` just ends up doing the same as `mypy myprogram.py`. Except instead of throwing away all the work used to parse your code, build the abstract syntax tree, etc. (and requiring you to then run `python myprogram.py` afterwards to actually use it), it can keep it around and just run your code immediately. - -TL;DR: `jax.jit` is amazing. diff --git a/README.md b/README.md index df14556..123a263 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,7 @@ Type annotations **and runtime checking** for: ```python from jaxtyping import Array, Float, PyTree +# Accepts floating-point 2D arrays with matching dimensions def matrix_multiply(x: Float[Array, "dim1 dim2"], y: Float[Array, "dim2 dim3"] ) -> Float[Array, "dim1 dim3"]: From 140be9ececbf2f5e484828fbc5815e485b51e31c Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Tue, 30 Aug 2022 13:26:41 -0700 Subject: [PATCH 4/9] Doc updates 2 --- API.md | 4 ++-- FAQ.md | 35 ++++++++++++++++++++++------------- 2 files changed, 24 insertions(+), 15 deletions(-) diff --git a/API.md b/API.md index 75f976a..a1b8496 100644 --- a/API.md +++ b/API.md @@ -18,7 +18,7 @@ In addition some modifiers can be applied: - 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[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#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. +When using multiple modifiers, their order does not matter. As a special case: - `...`: anonymous zero or more axes (equivalent to `*_`) e.g. `"... c h w"` @@ -60,7 +60,7 @@ f32[Array, "some_shape"] ### Array -The array should typically be a `jaxtyping.Array`, which is an alias for `jnp.ndarray`. +The array should typically be a `jaxtyping.Array`, which is an alias for `jax.numpy.ndarray`. But you can use other types as well. `jaxtyping` has support for JAX, NumPy, TensorFlow, and PyTorch, e.g.: ```python diff --git a/FAQ.md b/FAQ.md index b62f847..51f16f3 100644 --- a/FAQ.md +++ b/FAQ.md @@ -1,5 +1,17 @@ # FAQ +## Does jaxtyping work with static type checkers like `mypy`/`pyright`/`pytype`? + +There is partial support for these. An annotation of the form `dtype[array, shape]` should be treated as just `array` by a static type checker. Unfortunately full dtype/shape checking is beyond the scope of what static type checking is currently capable of. + +(Note that at time of writing, `pytype` has a bug in that `dtype[array, shape]` is sometimes treated as `Any` rather than `array`. `mypy` and `pyright` both work fine.) + +## How does jaxtyping interact with `jax.jit`? + +jaxtyping and `jax.jit` synergise beautifully. + +When calling JAX operations wrapped in a `jax.jit`, then the dtype/shape-checking will happen at trace time. (When JAX traces your function prior to compiling it.) The actual compiled code does not have any dtype/shape-checking, and will therefore still be just as fast as before! + ## `flake8` is throwing an error. In type annotations, strings are used for two different things. Sometimes they're strings. Sometimes they're "forward references", used to refer to a type that will be defined later. @@ -8,18 +20,6 @@ Some tooling in the Python ecosystem assumes that only the latter is true, and w In the case of `flake8`, at least, this is easily resolved. Multi-dimensional arrays (e.g. `f32[Array, "b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `f32[Array, "x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `f32[Array, " x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before. -## Does jaxtyping work with static type checkers like `mypy`/`pyright`/`pytype`? - -There is partial support for these. An annotation of the form `dtype[array, shape]` should be treated as just `array` by a static type checker. Unfortunately full dtype/shape checking is beyond the scope of what static type checking is currently capable of. - -(Note that at time of writing, `pytype` has a bug in that `dtype[array, shape]` is sometimes treated as `Any` rather than `array`. The other two work fine.) - -## How does jaxtyping interact with `jax.jit`? - -jaxtyping and `jax.jit` synergise beautifully. - -When calling JAX operations wrapped in a `jax.jit`, then the dtype/shape-checking will happen at trace time. (When JAX traces your function prior to compiling it.) The actual compiled code does not have any dtype/shape-checking, and will therefore still be just as fast as before! - ## Does jaxtyping use [PEP 646](https://www.python.org/dev/peps/pep-0646/) (variadic generics)? The intention of PEP 646 was to make it possible for static type checkers to perform shape checks of arrays. Unfortunately, this still isn't yet practical, so jaxtyping deliberately does not use this. (Yet?) @@ -36,4 +36,13 @@ The real problem is that Python's static typing ecosystem is a complicated colle 5. The syntax for static typing is verbose. You have to write things like `Array[Float32, Unpack[AnyShape], Literal[3], Height, Width]` instead of `f32[Array, "... 3 height width"]`. -6. [The underlying type system has flaws](https://github.com/patrick-kidger/torchtyping/issues/37#issuecomment-1153294196). [The numeric tower is broken](https://stackoverflow.com/a/69383462); [int is not a number](https://github.com/python/mypy/issues/3186#issuecomment-885718629); [virtual base classes don't work](https://github.com/python/mypy/issues/2922); [complex lies about having comparison operations, so type checkers have to lie about that lie in order to remove them again](https://posita.github.io/numerary/0.4/whytho/); `typing.*` don't work with `isinstance`; co/contra-variance are baked into containers (not specified at use-time); `dict` is variadic despite... not being variadic; bool is a subclass of int (!); ... etc. etc. +6. [The underlying type system has flaws](https://github.com/patrick-kidger/torchtyping/issues/37#issuecomment-1153294196). + [The numeric tower is broken](https://stackoverflow.com/a/69383462); + [int is not a number](https://github.com/python/mypy/issues/3186#issuecomment-885718629); + [virtual base classes don't work](https://github.com/python/mypy/issues/2922); + [complex lies about having comparison operations, so type checkers have to lie about that lie in order to remove them again](https://posita.github.io/numerary/0.4/whytho/); + `typing.*` don't work with `isinstance`; + co/contra-variance are baked into containers (not specified at use-time); + `dict` is variadic despite... not being variadic; + bool is a subclass of int (!); + ... etc. etc. From a53fe6af572ff917ab79bdfc900ae8d0f42e83d9 Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Wed, 7 Sep 2022 11:09:48 -0700 Subject: [PATCH 5/9] Changes from feedback: - Float32 introduced to replace f32 (etc.) - Old f32 aliases were previously around for backward-compatibility, but the code is pretty hideous, and all stakeholders are now on board through RFC #13. - Overall feedback was that Int{Sign,Unsign} was too long and not enough like numpy. Changed to Int, UInt, and Integer. --- API.md | 38 ++++---- jaxtyping/__init__.py | 43 ++++---- jaxtyping/array_types.py | 97 +++++++------------ test/import_hook_tester_beartype.py | 4 +- test/import_hook_tester_broken_checker.py | 4 +- .../another_file.py | 4 +- test/import_hook_tester_typeguard.py | 4 +- test/test_array.py | 80 ++++++++++----- 8 files changed, 138 insertions(+), 136 deletions(-) diff --git a/API.md b/API.md index a1b8496..58def13 100644 --- a/API.md +++ b/API.md @@ -38,14 +38,14 @@ The dtype should be any one of (imported from `jaxtyping`): - 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) + - Of particular precision: `BFloat16`, `Float16`, `Float32`, `Float64` - 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` + - Of particular precision: `Complex64`, `Complex128` + - Any integer or unsigned intger: `Integer` + - Any unsigned integer: `UInt` + - Of particular precision: `UInt8`, `UInt16`, `UInt32`, `UInt64` + - Any signed integer: `Int` + - Of particular precision: `Int8`, `Int16`, `Int32`, `Int64` 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 @@ -54,8 +54,8 @@ Float[Array, "some_shape"] ``` rather than ```python -from jaxtyping import Array, f32 -f32[Array, "some_shape"] +from jaxtyping import Array, Float32 +Float32[Array, "some_shape"] ``` ### Array @@ -93,7 +93,7 @@ Example: ```python # Import both the annotation and the `jaxtyped` decorator from `jaxtyping` -from jaxtyping import Array, f32, jaxtyped +from jaxtyping import Array, Float32, jaxtyped # Use your favourite typechecker: usually one of the two lines below. from typeguard import typechecked as typechecker @@ -102,9 +102,9 @@ from beartype import beartype as typechecker # Write your function. @jaxtyped must be applied above @typechecker! @jaxtyped @typechecker -def batch_outer_product(x: f32[Array, "b c1"], - y: f32[Array, "b c2"] - ) -> f32[Array, "b c1 c2"]: +def batch_outer_product(x: Float32[Array, "b c1"], + y: Float32[Array, "b c2"] + ) -> Float32[Array, "b c1 c2"]: return x[:, :, None] * y[:, None, :] ``` @@ -173,9 +173,9 @@ install_import_hook("do_stuff", ("typeguard", "typechecked")) import do_stuff ### do_stuff.py -from jaxtyping import Array, f32 +from jaxtyping import Array, Float32 -def g(x: f32[Array, "..."]): +def g(x: Float32[Array, "..."]): ... ``` @@ -204,17 +204,17 @@ which will apply the import hook to all modules whose names start with either `f The base class of all dtypes. This can be used to create your own custom collection of dtypes (analogous to `Float`, `Inexact` etc.) For example: ```python -class u8_or_u16(AbstractDtype): +class UInt8or16(AbstractDtype): dtypes = ["uint8", "uint16"] -u8_or_u16[Array, "shape"] +UInt8or16[Array, "shape"] ``` which is functionally equivalent to ```python -Union[u8[Array, "shape"], u16[Array, "shape"]] +Union[UInt8[Array, "shape"], UInt16[Array, "shape"]] ``` ### `jaxtyping.AbstractArray` The base class of all shape-and-dtype-specified arrays, e.g. it's a base class -for `f32[Array, "foo"]`. +for `Float32[Array, "foo"]`. diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index dfa38d2..413e8e1 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -17,43 +17,36 @@ # 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. +from jax.numpy import ndarray as Array + from .array_types import ( AbstractArray, AbstractDtype, - Array, - b, - bf16, + BFloat16, Bool, - c, - c64, - c128, Complex, - f, - f16, - f32, - f64, + Complex64, + Complex128, Float, + Float16, + Float32, + Float64, get_array_name_format, - i, - i8, - i16, - i32, - i64, Inexact, Int, - IntSign, - IntUnsign, - n, + Int8, + Int16, + Int32, + Int64, + Integer, Num, set_array_name_format, Shaped, - t, - u, - u8, - u16, - u32, - u64, - x, + UInt, + UInt8, + UInt16, + UInt32, + UInt64, ) from .decorator import jaxtyped from .import_hook import install_import_hook diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 7afbb37..c431171 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/array_types.py @@ -22,7 +22,6 @@ import functools as ft 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 @@ -251,19 +250,12 @@ class _MetaAbstractDtype(type): @ft.lru_cache(maxsize=None) 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 - if array_type is Array: - array_type = jnp.ndarray del item if not isinstance(dim_str, str): raise ValueError( @@ -409,7 +401,6 @@ class _MetaAbstractDtype(type): class AbstractDtype(metaclass=_MetaAbstractDtype): - deprecated: Optional[str] dtypes: Union[Literal[_any_dtype], List[str]] def __init__(self, *args, **kwargs): @@ -431,31 +422,29 @@ 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 BFloat16 from typing_extensions import Annotated as Bool from typing_extensions import Annotated as Complex + from typing_extensions import Annotated as Complex64 + from typing_extensions import Annotated as Complex128 from typing_extensions import Annotated as Float + from typing_extensions import Annotated as Float16 + from typing_extensions import Annotated as Float32 + from typing_extensions import Annotated as Float64 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 Int8 + from typing_extensions import Annotated as Int16 + from typing_extensions import Annotated as Int32 + from typing_extensions import Annotated as Int64 + from typing_extensions import Annotated as Integer 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 - - from jax.numpy import ndarray as Array + from typing_extensions import Annotated as UInt + from typing_extensions import Annotated as UInt8 + from typing_extensions import Annotated as UInt16 + from typing_extensions import Annotated as UInt32 + from typing_extensions import Annotated as UInt64 else: _bool = "bool" _uint8 = "uint8" @@ -473,30 +462,28 @@ else: _complex64 = "complex64" _complex128 = "complex128" - def _make_dtype(_dtypes, name, *, _deprecated=None): + def _make_dtype(_dtypes, name): class _Cls(AbstractDtype): - deprecated = _deprecated dtypes = _dtypes _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") + UInt8 = _make_dtype(_uint8, "u8") + UInt16 = _make_dtype(_uint16, "u16") + UInt32 = _make_dtype(_uint32, "u32") + UInt64 = _make_dtype(_uint64, "u64") + Int8 = _make_dtype(_int8, "i8") + Int16 = _make_dtype(_int16, "i16") + Int32 = _make_dtype(_int32, "i32") + Int64 = _make_dtype(_int64, "i64") + BFloat16 = _make_dtype(_bfloat16, "bf16") + Float16 = _make_dtype(_float16, "f16") + Float32 = _make_dtype(_float32, "f32") + Float64 = _make_dtype(_float64, "f64") + Complex64 = _make_dtype(_complex64, "c64") + Complex128 = _make_dtype(_complex128, "c128") uints = [_uint8, _uint16, _uint32, _uint64] ints = [_int8, _int16, _int32, _int64] @@ -506,25 +493,13 @@ else: # 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 - IntUnsign = _make_dtype(uints, "IntUnsign") - IntSign = _make_dtype(ints, "IntSign") - Int = _make_dtype(uints + ints, "Int") + Bool = _make_dtype(_bool, "Bool") + UInt = _make_dtype(uints, "UInt") + Int = _make_dtype(ints, "Int") + Integer = _make_dtype(uints + ints, "Integer") 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") - 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) - # Note that Array also has a non-deprecated use-case as - # `f32[Array, "foo"]`. - # TODO: once these deprecations are removed, then just have a - # `from jax.numpy import ndarray as Array` in `__init__.py`. - Array = _make_dtype(_bool, "Array", _deprecated=Shaped) + Shaped = _make_dtype(_any_dtype, "Shaped") diff --git a/test/import_hook_tester_beartype.py b/test/import_hook_tester_beartype.py index 11da8dd..414fa47 100644 --- a/test/import_hook_tester_beartype.py +++ b/test/import_hook_tester_beartype.py @@ -20,12 +20,12 @@ import jax.numpy as jnp import pytest -from jaxtyping import f32 +from jaxtyping import Float32 from .helpers import ParamError -def g(x: f32[jnp.ndarray, " b"]): +def g(x: Float32[jnp.ndarray, " b"]): pass diff --git a/test/import_hook_tester_broken_checker.py b/test/import_hook_tester_broken_checker.py index 11da8dd..414fa47 100644 --- a/test/import_hook_tester_broken_checker.py +++ b/test/import_hook_tester_broken_checker.py @@ -20,12 +20,12 @@ import jax.numpy as jnp import pytest -from jaxtyping import f32 +from jaxtyping import Float32 from .helpers import ParamError -def g(x: f32[jnp.ndarray, " b"]): +def g(x: Float32[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 9da8e5a..a7f710f 100644 --- a/test/import_hook_tester_transitive/another_file.py +++ b/test/import_hook_tester_transitive/another_file.py @@ -20,12 +20,12 @@ import jax.numpy as jnp import pytest -from jaxtyping import f32 +from jaxtyping import Float32 from ..helpers import ParamError -def g(x: f32[jnp.ndarray, " b"]): +def g(x: Float32[jnp.ndarray, " b"]): pass diff --git a/test/import_hook_tester_typeguard.py b/test/import_hook_tester_typeguard.py index 11da8dd..414fa47 100644 --- a/test/import_hook_tester_typeguard.py +++ b/test/import_hook_tester_typeguard.py @@ -20,12 +20,12 @@ import jax.numpy as jnp import pytest -from jaxtyping import f32 +from jaxtyping import Float32 from .helpers import ParamError -def g(x: f32[jnp.ndarray, " b"]): +def g(x: Float32[jnp.ndarray, " b"]): pass diff --git a/test/test_array.py b/test/test_array.py index bcdc4c0..c3c6082 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -21,7 +21,7 @@ import jax.numpy as jnp import jax.random as jr import pytest -from jaxtyping import Array, f32, Float, jaxtyped, Shaped +from jaxtyping import Array, Float, Float32, jaxtyped, Shaped from .helpers import ParamError, ReturnError @@ -35,6 +35,34 @@ def test_basic(typecheck): 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, + Int8, + Int16, + Int32, + Int64, + Num, + Shaped, + UInt, + UInt8, + UInt16, + UInt32, + UInt64, + ) + + def test_return(typecheck, getkey): @jaxtyped @typecheck @@ -92,12 +120,12 @@ def test_any_dtype(typecheck, getkey): def test_nested_jaxtyped(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[Array, "b c"], transpose: bool) -> f32[Array, "c b"]: + def g(x: Float32[Array, "b c"], transpose: bool) -> Float32[Array, "c b"]: return h(x, transpose) @jaxtyped @typecheck - def h(x: f32[Array, "c b"], transpose: bool) -> f32[Array, "b c"]: + def h(x: Float32[Array, "c b"], transpose: bool) -> Float32[Array, "b c"]: if transpose: return jnp.transpose(x) else: @@ -113,11 +141,11 @@ def test_nested_jaxtyped(typecheck, getkey): def test_nested_nojaxtyped(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[Array, "b c"]): + def g(x: Float32[Array, "b c"]): return h(x) @typecheck - def h(x: f32[Array, "c b"]): + def h(x: Float32[Array, "c b"]): return x with pytest.raises(ParamError): @@ -127,14 +155,14 @@ def test_nested_nojaxtyped(typecheck, getkey): def test_isinstance(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[Array, "b c"]) -> f32[Array, " z"]: + def g(x: Float32[Array, "b c"]) -> Float32[Array, " z"]: y = jnp.transpose(x) - assert isinstance(y, f32[Array, "c b"]) + assert isinstance(y, Float32[Array, "c b"]) assert not isinstance( - y, f32[Array, "b z"] + 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, f32[Array, "z"]) # z now bound + assert isinstance(out, Float32[Array, "z"]) # z now bound return out g(jr.normal(getkey(), (2, 3))) @@ -143,7 +171,9 @@ def test_isinstance(typecheck, getkey): def test_fixed(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[Array, "4 5 foo"], y: f32[Array, " foo"]) -> f32[Array, "4 5"]: + 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)) @@ -158,7 +188,7 @@ def test_fixed(typecheck, getkey): def test_anonymous(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[Array, "foo _"], y: f32[Array, " _"]): + def g(x: Float32[Array, "foo _"], y: Float32[Array, " _"]): pass a = jr.normal(getkey(), (3, 4)) @@ -169,7 +199,11 @@ def test_anonymous(typecheck, getkey): def test_named_variadic(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[Array, "*batch foo"], y: f32[Array, " *batch"], z: f32[Array, " foo"]): + def g( + x: Float32[Array, "*batch foo"], + y: Float32[Array, " *batch"], + z: Float32[Array, " foo"], + ): pass c = jr.normal(getkey(), (5,)) @@ -189,7 +223,7 @@ def test_named_variadic(typecheck, getkey): @jaxtyped @typecheck - def h(x: f32[Array, " foo *batch"], y: f32[Array, " foo *batch bar"]): + def h(x: Float32[Array, " foo *batch"], y: Float32[Array, " foo *batch bar"]): pass a = jr.normal(getkey(), (4,)) @@ -205,7 +239,7 @@ def test_named_variadic(typecheck, getkey): def test_anonymous_variadic(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[Array, "... foo"], y: f32[Array, " foo"]): + def g(x: Float32[Array, "... foo"], y: Float32[Array, " foo"]): pass a1 = jr.normal(getkey(), (5,)) @@ -227,7 +261,7 @@ def test_anonymous_variadic(typecheck, getkey): def test_broadcast_fixed(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[Array, "#4"]): + def g(x: Float32[Array, "#4"]): pass g(jr.normal(getkey(), (4,))) @@ -240,7 +274,7 @@ def test_broadcast_fixed(typecheck, getkey): def test_broadcast_named(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[Array, " #foo"], y: f32[Array, " #foo"]): + def g(x: Float32[Array, " #foo"], y: Float32[Array, " #foo"]): pass a = jr.normal(getkey(), (3,)) @@ -264,7 +298,7 @@ def test_broadcast_named(typecheck, getkey): def test_broadcast_variadic_named(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[Array, " *#foo"], y: f32[Array, " *#foo"]): + def g(x: Float32[Array, " *#foo"], y: Float32[Array, " *#foo"]): pass a = jr.normal(getkey(), (3,)) @@ -322,28 +356,28 @@ def test_broadcast_variadic_named(typecheck, getkey): def test_no_commas(): with pytest.raises(ValueError): - f32[Array, "foo, bar"] + Float32[Array, "foo, bar"] def test_symbolic(typecheck, getkey): @jaxtyped @typecheck - def make_slice(x: f32[Array, " dim"]) -> f32[Array, " dim-1"]: + def make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]: return x[1:] @jaxtyped @typecheck - def cat(x: f32[Array, " dim"]) -> f32[Array, " 2*dim"]: + def cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]: return jnp.concatenate([x, x]) @jaxtyped @typecheck - def bad_make_slice(x: f32[Array, " dim"]) -> f32[Array, " dim-1"]: + def bad_make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]: return x @jaxtyped @typecheck - def bad_cat(x: f32[Array, " dim"]) -> f32[Array, " 2*dim"]: + def bad_cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]: return jnp.concatenate([x, x, x]) x = jr.normal(getkey(), (5,)) @@ -365,7 +399,7 @@ def test_symbolic(typecheck, getkey): def test_incomplete_symbolic(typecheck, getkey): @jaxtyped @typecheck - def foo(x: f32[Array, " 2*dim"]): + def foo(x: Float32[Array, " 2*dim"]): pass x = jr.normal(getkey(), (4,)) From f28b0c789a8240c52cd01dcefc51d4577989adac Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Wed, 7 Sep 2022 11:32:14 -0700 Subject: [PATCH 6/9] doc tweaks --- API.md | 56 ++++++++++++++++++++++++++++++-------------------------- FAQ.md | 4 ++-- 2 files changed, 32 insertions(+), 28 deletions(-) diff --git a/API.md b/API.md index 58def13..ccb1210 100644 --- a/API.md +++ b/API.md @@ -73,7 +73,7 @@ Float[torch.Tensor, "..."] ### `jaxtyping.PyTree` -Each PyTree is denoted by a type `PyTree[LeafType]`, such as `PyTree[int]` or `PyTree[Union[str, f32[Array, "b c"]]]`. +Each PyTree is denoted by a type `PyTree[LeafType]`, such as `PyTree[int]` or `PyTree[Union[str, Float32[Array, "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)) @@ -85,7 +85,7 @@ To enable multi-argument consistency checks (i.e. that shapes match up between a Regardless of your choice, **this approach synergises beautifully with `jax.jit`!** All shape checks will be performed at trace-time only, and will not impact runtime performance. -### `jaxtyping.jaxtyped` +### Option 1: `jaxtyping.jaxtyped` Decorate a function with this to have shapes checked for consistency across multiple arguments. @@ -122,14 +122,14 @@ this function use the same axes sizes as the function it was called from. Likewise, this means you can use `isinstance` checks inside a function body and have them contribute to the same collection of consistency checks performed -by a typechecker against its arguments. (Or even forgo a typechecker altogether, -and just do your own manual `isinstance` checks.) +by a typechecker against its arguments. (Or even forgo a typechecker that analyses arguments, +and instead just do your own manual `isinstance` checks.) Only `isinstance` checks that pass will contribute to the store of axis name-size pairs; those that fail will not. As such it is safe to write e.g. `assert not isinstance(x, -f32[Array, "foo"])`. +Float32[Array, "foo"])`. -### `jaxtyping.install_import_hook` +### Option 2: `jaxtyping.install_import_hook` It can be a lot of effort to add `@jaxtyped` decorators all over your codebase. (Not to mention that double-decorators everywhere are a bit ugly.) The easier @@ -139,25 +139,24 @@ Example: ```python from jaxtyping import install_import_hook -# Plus either one of the following: -install_import_hook("foo", ("typeguard", "typechecked")) # decorate @jaxtyped and @typeguard.typechecked -install_import_hook("foo", ("beartype", "beartype")) # decorate @jaxtyped and @beartype.beartype -install_import_hook("foo", None) # decorate only @jaxtyped (if you have manually applied typechecking decorators) +# Plus any one of the following: + +# decorate @jaxtyped and @typeguard.typechecked +with install_import_hook("foo", ("typeguard", "typechecked")): + import foo + import foo.bar + import foo.bar.qux + +# decorate @jaxtyped and @beartype.beartype +with install_import_hook("foo", ("beartype", "beartype")): + ... + +# decorate only @jaxtyped (if you have manually applied typechecking decorators) +with install_import_hook("foo", None): + ... ``` -Any module imported **afterwards**, whose name begins with the specified string, will automatically have both `@jaxtyped` and the specified typechecker applied to all of their functions. (E.g. in the above example `foo`, `foo.bar`, `foo.bar.qux` would all be hook'd). - -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 -``` +Any module imported inside the `with` block, whose name begins with the specified string, will automatically have both `@jaxtyped` and the specified typechecker applied to all of their functions. (E.g. in the above example `foo`, `foo.bar`, `foo.bar.qux` would all be hook'd). The import hook can be applied to multiple packages via ```python @@ -169,8 +168,8 @@ install_import_hook(["foo", "bar.baz"], ...) ```python ### entry_point.py from jaxtyping import install_import_hook -install_import_hook("do_stuff", ("typeguard", "typechecked")) -import do_stuff +with install_import_hook("do_stuff", ("typeguard", "typechecked")): + import do_stuff ### do_stuff.py from jaxtyping import Array, Float32 @@ -187,7 +186,6 @@ from jaxtyping import install_import_hook 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 @@ -198,6 +196,12 @@ pytest --jaxtyping-packages=foo,bar.baz,beartype.beartype ``` which will apply the import hook to all modules whose names start with either `foo` or `bar.baz`. The typechecker used in this example is `beartype.beartype`. +## Static type checking + +jaxtyping should be compatible with static type checkers (the big three are `mypy`, `pyright`, `pytype`) out of the box. + +Due to limitations of static type checkers, only the array type (JAX array vs NumPy array vs PyTorch tensor vs TensorFlow tensor) is checked. Shape and dtype are not checked. [See the FAQ](./FAQ.md#what-about-pep-646-and-variadic-generics) for more details. + ## Abstract base classes ### `jaxtyping.AbstractDtype` diff --git a/FAQ.md b/FAQ.md index 51f16f3..7291d5b 100644 --- a/FAQ.md +++ b/FAQ.md @@ -18,7 +18,7 @@ In type annotations, strings are used for two different things. Sometimes they'r Some tooling in the Python ecosystem assumes that only the latter is true, and will throw spurious errors if you try to use a string just as a string (like we do). -In the case of `flake8`, at least, this is easily resolved. Multi-dimensional arrays (e.g. `f32[Array, "b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `f32[Array, "x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `f32[Array, " x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before. +In the case of `flake8`, at least, this is easily resolved. Multi-dimensional arrays (e.g. `Float32[Array, "b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `Float32[Array, "x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `Float32[Array, " x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before. ## Does jaxtyping use [PEP 646](https://www.python.org/dev/peps/pep-0646/) (variadic generics)? @@ -34,7 +34,7 @@ The real problem is that Python's static typing ecosystem is a complicated colle 4. Variadic generics exist. Variadic protocols do not. (It's not clear that these were contemplated.) -5. The syntax for static typing is verbose. You have to write things like `Array[Float32, Unpack[AnyShape], Literal[3], Height, Width]` instead of `f32[Array, "... 3 height width"]`. +5. The syntax for static typing is a little verbose. You have to write things like `Array[Float32, Unpack[AnyShape], Literal[3], Height, Width]` instead of `Float32[Array, "... 3 height width"]`. 6. [The underlying type system has flaws](https://github.com/patrick-kidger/torchtyping/issues/37#issuecomment-1153294196). [The numeric tower is broken](https://stackoverflow.com/a/69383462); From c82fbbea4eef2a4cda137eeaa4c21e7c060452cc Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Wed, 7 Sep 2022 12:00:58 -0700 Subject: [PATCH 7/9] doc tweaks --- API.md | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/API.md b/API.md index ccb1210..b0a6e36 100644 --- a/API.md +++ b/API.md @@ -6,7 +6,7 @@ Each array is denoted by a type `dtype[array, shape]`, such as `Float[Array, "ba ### Shape -The shape should be a string of space-separated symbols, such as "a b c d". Each symbol can be either an: +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. `"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[Array, "dim"]) -> Float[Array, "dim-1"]`. @@ -15,7 +15,7 @@ When calling a function, variable-size axes and symbolic axes will be matched up In addition some modifiers can be applied: - 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[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#foo"]`. +- Prepend `#` to a dimension to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable, e.g. `def add(x: Float[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#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 _ _"`. When using multiple modifiers, their order does not matter. @@ -27,7 +27,7 @@ Some notes: - To denote a scalar shape use `""`, e.g. `Float[Array, ""]`. - To denote an arbitrary shape (and only check dtype) use `"..."`, e.g. `Float[Array, "..."]`. - 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: Float[Array, "*#foo"], y: Float[Array, "*#foo"]) -> Float[Array, "*#foo"]`. +- An example of broadcasting multiple dimensions: `def add(x: Float[Array, "*#foo"], y: Float[Array, "*#foo"]) -> Float[Array, "*#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 @@ -132,31 +132,36 @@ Float32[Array, "foo"])`. ### Option 2: `jaxtyping.install_import_hook` It can be a lot of effort to add `@jaxtyped` decorators all over your codebase. -(Not to mention that double-decorators everywhere are a bit ugly.) The easier -option is usually to use the import import hook. +(Not to mention that double-decorators everywhere are a bit ugly.) -Example: +The easier option is usually to use the import hook. +This can be used via a `with` block; for example: ```python from jaxtyping import install_import_hook # Plus any one of the following: # decorate @jaxtyped and @typeguard.typechecked with install_import_hook("foo", ("typeguard", "typechecked")): - import foo - import foo.bar - import foo.bar.qux + import foo # Any module imported inside this `with` block, whose name begins + import foo.bar # with the specified string, will automatically have both `@jaxtyped` + import foo.bar.qux # and the specified typechecker applied to all of their functions. # decorate @jaxtyped and @beartype.beartype with install_import_hook("foo", ("beartype", "beartype")): ... -# decorate only @jaxtyped (if you have manually applied typechecking decorators) +# decorate only @jaxtyped (if you want that for some reason) with install_import_hook("foo", None): ... ``` -Any module imported inside the `with` block, whose name begins with the specified string, will automatically have both `@jaxtyped` and the specified typechecker applied to all of their functions. (E.g. in the above example `foo`, `foo.bar`, `foo.bar.qux` would all be hook'd). +If you don't like using the `with` block, the hook can be used without that: +```python +hook = install_import_hook(...): +import ... +hook.uninstall() +``` The import hook can be applied to multiple packages via ```python From cb11a93b2285de8bfb84f5a0b82480f7a3195c95 Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Wed, 7 Sep 2022 12:06:15 -0700 Subject: [PATCH 8/9] doc tweaks --- API.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/API.md b/API.md index b0a6e36..c1ad912 100644 --- a/API.md +++ b/API.md @@ -205,7 +205,7 @@ which will apply the import hook to all modules whose names start with either `f jaxtyping should be compatible with static type checkers (the big three are `mypy`, `pyright`, `pytype`) out of the box. -Due to limitations of static type checkers, only the array type (JAX array vs NumPy array vs PyTorch tensor vs TensorFlow tensor) is checked. Shape and dtype are not checked. [See the FAQ](./FAQ.md#what-about-pep-646-and-variadic-generics) for more details. +Due to limitations of static type checkers, only the array type (JAX array vs NumPy array vs PyTorch tensor vs TensorFlow tensor) is checked. Shape and dtype are not checked. [See the FAQ](./FAQ.md) for more details. ## Abstract base classes From 01f8f20bf5962bb48ae374f1a0149db63e20408c Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Wed, 7 Sep 2022 12:20:42 -0700 Subject: [PATCH 9/9] Fixed a few names --- jaxtyping/array_types.py | 36 ++++++++++++++++++------------------ test/test_array.py | 6 +++++- 2 files changed, 23 insertions(+), 19 deletions(-) diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index c431171..631c281 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/array_types.py @@ -253,7 +253,7 @@ class _MetaAbstractDtype(type): 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. For example `jaxtyping.Float32[jnp.ndarray, 'foo bar']`." ) array_type, dim_str = item del item @@ -406,7 +406,7 @@ class AbstractDtype(metaclass=_MetaAbstractDtype): def __init__(self, *args, **kwargs): raise RuntimeError( "AbstractDtype cannot be instantiated. Perhaps you wrote e.g. " - '`f32("shape")` when you mean `f32[jnp.ndarray, "shape"]`?' + '`Float32("shape")` when you mean `Float32[jnp.ndarray, "shape"]`?' ) def __init_subclass__(cls, **kwargs): @@ -470,20 +470,20 @@ else: _Cls.__qualname__ = name return _Cls - UInt8 = _make_dtype(_uint8, "u8") - UInt16 = _make_dtype(_uint16, "u16") - UInt32 = _make_dtype(_uint32, "u32") - UInt64 = _make_dtype(_uint64, "u64") - Int8 = _make_dtype(_int8, "i8") - Int16 = _make_dtype(_int16, "i16") - Int32 = _make_dtype(_int32, "i32") - Int64 = _make_dtype(_int64, "i64") - BFloat16 = _make_dtype(_bfloat16, "bf16") - Float16 = _make_dtype(_float16, "f16") - Float32 = _make_dtype(_float32, "f32") - Float64 = _make_dtype(_float64, "f64") - Complex64 = _make_dtype(_complex64, "c64") - Complex128 = _make_dtype(_complex128, "c128") + UInt8 = _make_dtype(_uint8, "UInt8") + UInt16 = _make_dtype(_uint16, "UInt16") + UInt32 = _make_dtype(_uint32, "UInt32") + UInt64 = _make_dtype(_uint64, "UInt64") + Int8 = _make_dtype(_int8, "Int8") + Int16 = _make_dtype(_int16, "Int16") + Int32 = _make_dtype(_int32, "Int32") + Int64 = _make_dtype(_int64, "Int64") + BFloat16 = _make_dtype(_bfloat16, "BFloat16") + Float16 = _make_dtype(_float16, "Float16") + Float32 = _make_dtype(_float32, "Float32") + Float64 = _make_dtype(_float64, "Float64") + Complex64 = _make_dtype(_complex64, "Complex64") + Complex128 = _make_dtype(_complex128, "Complex128") uints = [_uint8, _uint16, _uint32, _uint64] ints = [_int8, _int16, _int32, _int64] @@ -499,7 +499,7 @@ else: Integer = _make_dtype(uints + ints, "Integer") 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 + Inexact = _make_dtype(floats + complexes, "Inexact") + Num = _make_dtype(uints + ints + floats + complexes, "Num") Shaped = _make_dtype(_any_dtype, "Shaped") diff --git a/test/test_array.py b/test/test_array.py index c3c6082..eb377b3 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -21,7 +21,7 @@ import jax.numpy as jnp import jax.random as jr import pytest -from jaxtyping import Array, Float, Float32, jaxtyped, Shaped +from jaxtyping import AbstractDtype, Array, Float, Float32, jaxtyped, Shaped from .helpers import ParamError, ReturnError @@ -62,6 +62,10 @@ def test_dtypes(): UInt64, ) + for key, val in locals().items(): + if issubclass(val, AbstractDtype): + assert key == val.__name__ + def test_return(typecheck, getkey): @jaxtyped