Jake VanderPlas
5cc689976f
Use PEP484-style exports in several submodules
2024-08-14 08:59:56 -07:00
Peter Hawkins
ba557d5e1b
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
...
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.
PiperOrigin-RevId: 476167538
2022-09-22 12:27:19 -07:00
Jake VanderPlas
5782210174
CI: fix flake8 ignore declarations
2022-04-21 13:44:12 -07:00
Roy Frostig
30c3a39467
introduce custom_batching.sequential_vmap
...
An anticipated common use of `custom_vmap` is in order to implement a
map via loop (i.e. to sequentially apply the mapped function), instead
of actually vectorizing.
2022-01-20 21:37:58 -08:00
Roy Frostig
1764d5a3fc
introduce custom_vmap and a corresponding custom_vmap_p primitive
...
Adds an impl and a batching rule for the primitive.
Co-authored-by: Matthew Johnson <mattjj@google.com>
2022-01-05 18:07:20 -08:00
