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,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"]`.

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

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")

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
 
 

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
 
 

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
 
 

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
 
 

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