Rewrote syntax

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 <ins>n</ins>umber)
-    - Any floating or complex: `x` (for ine<ins>x</ins>act)
-      - 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 in<ins>t</ins>eger)
-      - 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"]`.

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
 

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"

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)

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
 
 

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
 
 

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
 
 

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
 
 

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

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