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] 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,))