Yash Katariya 4db212d2c6 Add _sharding argument to broadcasted_iota as a private parameter which only works under sharding_in_types mode. ...

This is required because `jax.nn.one_hot` calls into `broascasted_iota`.

PiperOrigin-RevId: 687152343

2024-10-17 21:16:51 -07:00

..

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

functions.py

Add _sharding argument to broadcasted_iota as a private parameter which only works under sharding_in_types mode.

2024-10-17 21:16:51 -07:00

initializers.py

update README and several docs to typed RNG keys

2024-08-11 08:09:47 -07:00