mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src/lax
History
Praveen Narayanan b6d4fe5387 Define lax.ragged_dot_general and express lax.ragged_dot in terms of it. 
PiperOrigin-RevId: 735471245
2025-03-10 12:25:22 -07:00
..
control_flow
[scan] when num_trips==0, don't generate weird size-zero reshapes
2025-03-07 21:35:40 +00:00
__init__.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
ann.py
Clean up and fix primal type to tangent type mapping
2024-09-18 13:43:54 -07:00
convolution.py
Clone RandomAlgorithm into lax.py, instead of using the version from XLA.
2024-10-11 11:03:15 -07:00
eigh.py
refactor: import numpy objects directly in jax.numpy
2025-02-14 12:47:58 -08:00
fft.py
Clone RandomAlgorithm into lax.py, instead of using the version from XLA.
2024-10-11 11:03:15 -07:00
lax.py
Define lax.ragged_dot_general and express lax.ragged_dot in terms of it.
2025-03-10 12:25:22 -07:00
linalg.py
Move (most) jaxlib linalg custom call registration into JAX.
2025-03-10 08:17:44 -07:00
other.py
Simplify and consolidate dot algorithm control in lax.
2024-10-07 13:21:34 -07:00
parallel.py
Now that sharding_in_types config flag is True, remove the config and all the conditionals
2025-02-19 06:53:35 -08:00
qdwh.py
Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708
2024-09-03 09:49:22 -07:00
slicing.py
Add a divisibility check so that we make sure that sharding evenly divides the shape (until this restriction is lifted) to make sure we don't create bad shardings.
2025-03-07 07:01:34 -08:00
special.py
Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.
2023-11-17 11:47:14 -08:00
stack.py
Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708
2024-09-03 09:49:22 -07:00
svd.py
Enable pivoted QR on CPU devices.
2025-01-09 20:44:45 +00:00
utils.py
Start adding primitive registration helper functions to lax.linalg.
2025-02-21 04:05:34 -08:00
windowed_reductions.py
[sharding_in_types] Add sharding rules for the following primitives:
2025-02-22 10:45:58 -08:00