Roy Frostig

RNGs: key types and custom implementations

4

This issue tracks the introduction of typed random key arrays with pluggable PRNG implementations. Work on this began with #6899, which mostly works out the "pluggable implementations" part, makes initial...

custom batching (vmap)

1

Support custom batching, i.e. the ability to register a custom "vmap rule" for any given function. Example usage would look something like: ```python from jax import vmap, custom_vmap, numpy as...

efficient untrue batching of `random_bit_generator`

1

The batching rule for the `random_bit_generator` primitive, over a batch of keys, emits a loop (via `lax.map`): https://github.com/google/jax/blob/42ae8432185bf03f61ddd2e7bc279d3abb5247fd/jax/_src/lax/control_flow/loops.py#L2012-L2024 This is a workaround to the corresponding `RandomBitGenerator` HLO not being batchable....

extensions module

3

Several projects and libraries critically rely on jax internals, such the Jaxpr IR. [Oryx](https://github.com/jax-ml/oryx) is one example among many others. These projects use the jax core as a sort of...

shape hint arguments in lax numpy functions to enable compilation

1

Several functions in `jax.numpy` are not `jit`-compilable for shape-level arguments. Specifically, we can't decide their output shape from input shapes alone. We've given [`jax.numpy.bincount`](https://github.com/google/jax/blob/db8f66d508d536e1312291ba966905ddf910582d/jax/numpy/lax_numpy.py#L1309) an optional `length` argument that's not...

### Description ```python import jax def f(x): return x + jax.random.randint(jax.random.key(72), (), 0, 10) def g(x): with jax.threefry_partitionable(True): # False by default return x + jax.random.randint(jax.random.key(72), (), 0, 10) h...

### Description Consider the interaction of the following features: 1. `jax.remat` 2. `rbg` or `unsafe_rbg` RNG implementations 3. the XLA flag `xla_tpu_spmd_rng_bit_generator_unsafe` 4. automatic partitioning `remat` says: jax, you are...

This is a batch of individually minor edits that aim for easier/clearer reading and some decluttering of titles and headings (which hopefully also helps in indices and in our navigation...