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

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