Ayaka c60bafcc33 [Pallas TPU] Fix lowering for jnp.remainder ...

Fixes https://github.com/jax-ml/jax/issues/24027

PiperOrigin-RevId: 688614799

2024-10-22 11:01:58 -07:00

..

BUILD

Move additional CI enabled/disabled configurations into jax BUILD files.

2024-10-10 08:41:45 -07:00

export_back_compat_pallas_test.py

[Pallas] Update export compatibility tests

2024-09-24 09:17:56 -07:00

export_pallas_test.py

[pallas] Add lowering errors for block shapes that are not supported.

2024-07-29 06:49:27 -07:00

gpu_attention_test.py

[Pallas] Fully skip GPU attention tests on win32.

2024-09-09 10:21:33 -07:00

gpu_ops_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

indexing_test.py

[Pallas] Fix array indexing error when dimension size is not a multiple of stride

2024-09-10 02:52:55 +01:00

mosaic_gpu_test.py

[Pallas:MGPU] Add support for grid dims in GPUMesh

2024-10-22 04:10:46 -07:00

ops_test.py

[Pallas TPU] Fix lowering for jnp.remainder

2024-10-22 11:01:58 -07:00

pallas_error_handling_test.py

[XLA:Mosaic][Pallas] Enable vector.ExtractOp for non-zero indices.

2024-09-26 13:57:45 -07:00

pallas_jumble_test.py

Add a memory saving index rewrite step to vmap with ragged inputs over pallas_call.

2024-10-14 14:01:42 -07:00

pallas_shape_poly_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

2024-07-25 01:50:36 -07:00

pallas_test.py

[Pallas TPU] Add a better error message for rank 1 block mappings check.

2024-10-16 08:41:22 -07:00

pallas_vmap_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

tpu_all_gather_test.py

[pallas] Renamed platform-specific tests.

2024-07-10 08:23:06 -07:00

tpu_gmm_test.py

[pallas] Renamed platform-specific tests.

2024-07-10 08:23:06 -07:00

tpu_ops_test.py

[Mosaic TPU] The previous change does not actually force the input offsets read by the rules, but simply disables all the checks. Reverting so that we at least regain the checks until we have a proper fix.

2024-10-01 02:23:52 -07:00

tpu_paged_attention_kernel_test.py

Skip some Pallas tests that fail on TPUv6.

2024-07-23 07:16:24 -07:00

tpu_pallas_async_test.py

[Pallas] Add state discharge rule for pallas_call

2024-10-01 16:30:56 -07:00

tpu_pallas_distributed_test.py

Do not allow None axis names in meshes.

2024-10-16 10:32:25 -07:00

tpu_pallas_pipeline_test.py

Temporarily disable broken tests in tpu_pallas_pipeline_test.py

2024-08-08 14:04:04 -07:00

tpu_pallas_random_test.py

[Pallas] Add lowering for threefry PRNG.

2024-10-09 14:48:26 -07:00

tpu_pallas_state_test.py

[Pallas] Use core_map instead of shard_map for Shmallas

2024-10-15 03:26:58 -07:00

tpu_pallas_test.py

[pallas] run_scoped now supports partial discharge.

2024-10-18 10:22:31 -07:00

tpu_splash_attention_kernel_test.py

[Pallas/TPU] Fix bug with LocalMask grid shrinking

2024-07-25 04:05:57 -07:00

tpu_splash_attention_mask_test.py

[pallas] Renamed platform-specific tests.

2024-07-10 08:23:06 -07:00