mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src/lax
History
Peter Hawkins 0150d15cb2 Increase minimum jaxlib version to 0.3.7. 
Drop backwards compatibility with older jaxlib versions.
2022-04-18 08:09:50 -04:00
..
__init__.py
register the lax module for stack trace filtering
2021-02-22 21:11:41 -08:00
ann.py
[JAX] Fix batch logic for approx_min/max_k
2022-04-08 13:50:36 -07:00
control_flow.py
Switch gpu, cpu and jax2tf to use the new OptimizationBarrier op.
2022-04-13 22:51:36 -07:00
convolution.py
[MLIR] Change signature of lowering rules.
2021-12-15 19:06:58 -08:00
fft.py
Increase minimum jaxlib version to 0.3.7.
2022-04-18 08:09:50 -04:00
lax.py
Increase minimum jaxlib version to 0.3.7.
2022-04-18 08:09:50 -04:00
linalg.py
Increase minimum jaxlib version to 0.3.7.
2022-04-18 08:09:50 -04:00
other.py
Use real Precision type for lax.PrecisionType
2022-03-04 04:21:25 -08:00
parallel.py
[MHLO] Separate registrations for collective and initial_style primitives from the XLA translation rule registration.
2022-04-13 07:26:26 -07:00
polar.py
Document jax.lax.Precision
2021-11-08 14:15:31 -08:00
qdwh.py
[linalg] Fix type promotion in QDWH.
2022-03-03 20:59:28 -08:00
slicing.py
add some simple iree tests
2022-04-14 10:55:00 -07:00
svd.py
[JAX] Added jit-able singular value decomposition.
2022-02-03 11:16:55 -08:00
utils.py
[MHLO] Add explicit XLA translation rules for primitives that lack MHLO lowerings that rely on standard_primitive registering a translation rule.
2022-04-16 07:01:19 -07:00
windowed_reductions.py
fixup from 5415306: remove extraneous lines
2022-03-11 15:19:10 -08:00