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

Updated to support both new and old style JAX PRNG keys, as they are going to co-exist simultaneously. See https://github.com/google/jax/pull/17297

Browse Source
This commit is contained in:
Patrick Kidger
2023-09-14 20:00:46 -07:00
parent e347c480d5
commit e308695293
2 changed files with 8 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/api/array.md
+3 -4
View File
@@ -95,16 +95,15 @@ BatchImage = Float[Array, "batch channels height width"]
Note that `jaxtyping.{Array, ArrayLike}` are only available if JAX has been installed.
## Scalars, PRNGKeys
## Scalars, PRNG keys
For convenience, jaxtyping also includes `jaxtyping.Scalar`, `jaxtyping.ScalarLike`, and `jaxtyping.PRNGKeyArray`, defined as:
```python
Scalar = Shaped[Array, ""]
ScalarLike = Shaped[ArrayLike, ""]
# Depending on the value of `JAX_ENABLE_CUSTOM_PRNG`:
PRNGKeyArray = Key[Array, ""]
PRNGKeyArray = UInt32[Array, "2"]
# Left: new-style typed keys; right: old-style keys. See JEP 9263.
PRNGKeyArray = Union[Key[Array, ""], UInt32[Array, "2"]]
```
Recalling that shape-and-dtype specified jaxtyping arrays can be nested, this means that e.g. you can annotate the output of `jax.random.split` with `Shaped[PRNGKeyArray, "2"]`, or e.g. an integer scalar with `Int[Scalar, ""]`.
jaxtyping/_array_types.py
+5 -7
View File
@@ -657,12 +657,10 @@ Num = _make_dtype(uints + ints + floats + complexes, "Num")
Shaped = _make_dtype(_any_dtype, "Shaped")
if has_jax:
if jax.config.jax_enable_custom_prng:
_key_regex = re.compile(r"^key<\w+>$")
Key = _make_dtype(_key_regex, "Key")
PRNGKeyArray = Key[jax.Array, ""]
else:
Key = UInt32
PRNGKeyArray = Key[jax.Array, "2"]
_key_regex = re.compile(r"^key<\w+>$")
Key = _make_dtype(_key_regex, "Key")
# New-style `jax.random.key` have scalar shape and dtype `key<foo>`.
# Old-style `jax.random.PRNGKey` have shape `(2,)` and dtype `uint32`.
PRNGKeyArray = Union[Key[jax.Array, ""], UInt32[jax.Array, "2"]]
Scalar = Shaped[jax.Array, ""]
ScalarLike = Shaped[jax.typing.ArrayLike, ""]