wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-09 11:24:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

Standardised terminology: now using just "axis"/"axes", not "dimension"

Browse Source
This commit is contained in:
Patrick Kidger
2023-11-27 09:50:02 -08:00
parent 12d540794f
commit ba3b2027cc
4 changed files with 23 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -1
View File
@@ -10,7 +10,7 @@ Type annotations **and runtime type-checking** for:
```python
from jaxtyping import Array, Float, PyTree
# Accepts floating-point 2D arrays with matching dimensions
# Accepts floating-point 2D arrays with matching axes
def matrix_multiply(x: Float[Array, "dim1 dim2"],
y: Float[Array, "dim2 dim3"]
) -> Float[Array, "dim1 dim3"]:
docs/api/array.md
+5 -5
View File
@@ -19,12 +19,12 @@ 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.
- Prepend `*` to an axis to indicate that it can match multiple axes, e.g. `"*batch c h w"` will match zero or more batch axes.
- Prepend `#` to an axis 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 _ _"`.
- Prepend `_` to an axis to disable any runtime checking of that axis (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `"b c _ _"`.
- Documentation-only names (i.e. they're ignored by jaxtyping) can be handled by prepending a name followed by `=` e.g. `Float[Array, "rows=4 cols=3"]`.
- Prepend `?` to a dimension to indicate that its size can vary within a PyTree structure. (See [PyTree annotations](../pytree/).)
- Prepend `?` to an axis to indicate that its size can vary within a PyTree structure. (See [PyTree annotations](../pytree/).)
When using multiple modifiers, their order does not matter.
@@ -37,7 +37,7 @@ As a special case:
- 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:
- An example of broadcasting multiple axes:
`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.
docs/index.md
+1 -1
View File
@@ -24,7 +24,7 @@ The annotations provided by jaxtyping are compatible with runtime type-checking
```python
from jaxtyping import Array, Float, PyTree
# Accepts floating-point 2D arrays with matching dimensions
# Accepts floating-point 2D arrays with matching axes
def matrix_multiply(x: Float[Array, "dim1 dim2"],
y: Float[Array, "dim2 dim3"]
) -> Float[Array, "dim1 dim3"]:
jaxtyping/_array_types.py
+16 -16
View File
@@ -127,10 +127,10 @@ def _check_dims(
except NameError as e:
jaxtyping_raise_from(
NameError(
f"Cannot process symbolic dimension '{cls_dim.elem_string}' as "
"some dimension names have not been processed. In practice you "
"should usually only use symbolic dimensions in annotations "
"for return types, referring only to dimensions annotated for "
f"Cannot process symbolic axis '{cls_dim.elem_string}' as "
"some axis names have not been processed. In practice you "
"should usually only use symbolic axes in annotations "
"for return types, referring only to axes annotated for "
"arguments."
),
e,
@@ -341,18 +341,18 @@ def _make_array(array_type, dim_str, dtypes, name):
if "," in elem and "(" not in elem:
# Common mistake.
# Disable in the case that there's brackets to allow for function calls,
# e.g. `min(foo,bar)`, in symbolic dimensions.
raise ValueError("Dimensions should be separated with spaces, not commas")
# e.g. `min(foo,bar)`, in symbolic axes.
raise ValueError("Axes should be separated with spaces, not commas")
if elem.endswith("#"):
raise ValueError(
"As of jaxtyping v0.1.0, broadcastable dimensions are now denoted "
"As of jaxtyping v0.1.0, broadcastable axes are now denoted "
"with a # at the start, rather than at the end"
)
if "..." in elem:
if elem != "...":
raise ValueError(
"Anonymous multiple dimension '...' must be used on its own; "
"Anonymous multiple axes '...' must be used on its own; "
f"got {elem}"
)
broadcastable = False
@@ -382,7 +382,7 @@ def _make_array(array_type, dim_str, dtypes, name):
if variadic:
raise ValueError(
"Do not use * twice to denote accepting multiple "
"dimensions, e.g. `**foo` is not allowed"
"axes, e.g. `**foo` is not allowed"
)
variadic = True
elem = elem[1:]
@@ -421,7 +421,7 @@ def _make_array(array_type, dim_str, dtypes, name):
if variadic:
if index_variadic is not None:
raise ValueError(
"Cannot use multiple-dimension specifiers (`*name` or `...`) "
"Cannot use variadic specifiers (`*name` or `...`) "
"more than once."
)
index_variadic = index
@@ -429,7 +429,7 @@ def _make_array(array_type, dim_str, dtypes, name):
if dim_type is _DimType.fixed:
if variadic:
raise ValueError(
"Cannot have a fixed axis bind to multiple dimensions, e.g. "
"Cannot have a fixed axis bind to multiple axes, e.g. "
"`*4` is not allowed."
)
if anonymous:
@@ -446,7 +446,7 @@ def _make_array(array_type, dim_str, dtypes, name):
if anonymous:
if broadcastable:
raise ValueError(
"Cannot have a dimension be both anonymous and "
"Cannot have an axis be both anonymous and "
"broadcastable, e.g. `#_` is not allowed."
)
if variadic:
@@ -462,17 +462,17 @@ def _make_array(array_type, dim_str, dtypes, name):
assert dim_type is _DimType.symbolic
if anonymous:
raise ValueError(
"Cannot have a symbolic dimension be anonymous, e.g. "
"Cannot have a symbolic axis be anonymous, e.g. "
"`_foo+bar` is not allowed"
)
if variadic:
raise ValueError(
"Cannot have symbolic multiple-dimensions, e.g. "
"Cannot have symbolic multiple-axes, e.g. "
"`*foo+bar` is not allowed"
)
if treepath:
raise ValueError(
"Cannot have a symbolic dimensions with tree-path dependence, e.g. "
"Cannot have a symbolic axis with tree-path dependence, e.g. "
"`?foo+bar` is not allowed"
)
elem_string = elem
@@ -520,7 +520,7 @@ def _make_array(array_type, dim_str, dtypes, name):
index_variadic = array_type.index_variadic + len(dims)
else:
raise ValueError(
"Cannot use multiple-dimension specifiers (`*name` or `...`) "
"Cannot use variadic specifiers (`*name` or `...`) "
"in both the original array and the extended array"
)
dims = dims + array_type.dims