mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 22:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/interpreters
History
Peter Hawkins 3e78a0e290 Keep ShapedArray avals on xla.DeviceArray values 
Makes abstractification of DeviceArray values cheaper, which is on the critical path for executing a compiled function.
2019-08-12 10:03:04 -04:00
..
__init__.py
populating source tree
2018-11-17 18:03:33 -08:00
ad.py
Merge branch 'master' into custom-transforms
2019-07-03 07:51:05 +01:00
batching.py
Add a common lax._canonicalize_shape method, use on methods that accept shapes in lax.
2019-07-23 16:19:02 -04:00
parallel.py
Two more pmap optimizations:
2019-08-06 09:14:20 -04:00
partial_eval.py
JaxprTrace.process_map wasn't using env_tracers
2019-07-08 18:38:30 -07:00
pxla.py
Keep ShapedArray avals on xla.DeviceArray values
2019-08-12 10:03:04 -04:00
xla.py
Keep ShapedArray avals on xla.DeviceArray values
2019-08-12 10:03:04 -04:00