Peter Hawkins 6e6f693e6d Use lax.broadcasted_iota in jax.nn.one_hot. ...

Minor cleanup that means we emit one fewer MHLO op, no functional changes intended.

2022-04-20 09:09:19 -04:00

..

__init__.py

Move jax.nn implementation into jax._src.nn.

2020-10-17 13:45:01 -04:00

functions.py

Use lax.broadcasted_iota in jax.nn.one_hot.

2022-04-20 09:09:19 -04:00

initializers.py

Skip doctests for initializer examples.

2022-03-08 08:29:43 -05:00