ROCm/jax

Pull Requests Commits

[sharding_in_types] Aval sharding under full auto mode should contain None and not UNCONSTRAINED because axis_types + pspec give the full picture.

yashk2810 committed 360 days ago

755d6cda

[Pallas] Improvements to core_map

sharadmv committed 360 days ago

7be127f2

Increase the minimum SciPy version to 1.11.1.

hawkinsp committed 360 days ago

392a8517

[array api] update test suite to latest commit

jakevdp committed 361 days ago

f6c9e87d

Merge pull request #25042 from dfm:ffi-example-input-output-alias

Google-ML-Automation committed 361 days ago

f1777d5b

Sort manual axes when lowering `jax.shard_map` to `sdy.manual_computation`, which ensures the determinism in the generated `sdy.manual_computation`.

ZixuanJiang committed 361 days ago

64c0f62e

Add an example demonstrating input-output aliasing with the FFI.

dfm committed 361 days ago

62656b32

Update XLA dependency to use revision

Google-ML-Automation committed 361 days ago

00c363e1

[Pallas] Fix pallas_call lowering mutating compiler params during Triton lowering.

justinjfu committed 361 days ago

8c9a5394

[AutoPGLE] FIx PGLE kokoro test failures.

Google-ML-Automation committed 361 days ago

40238105

Merge pull request #25441 from Exferro:fixed_advanced_autodiff_doc

Google-ML-Automation committed 361 days ago

57c2afe7

Clean up `backend_or_name` vs. `platforms` in lowering code.

gnecula committed 361 days ago

fdb6af82

[Mosaic GPU] Use `num_q_heads=2` in `flash_attention.py`

andportnoy committed 361 days ago

Verified 5b808921

Add a register_custom_type_id function to the GPU plugins.

dfm committed 361 days ago

a7f384cc

Update the advanced autodiff tutorial and replace some vmap with grad

Exferro committed 361 days ago

f881f507

Merge pull request #25748 from shoyer:divmod

Google-ML-Automation committed 361 days ago

853af560

Merge pull request #25633 from dfm:move-ffi

Google-ML-Automation committed 361 days ago

56f0f953

Merge pull request #25731 from gnecula:poly_random_even

Google-ML-Automation committed 361 days ago

712bece2

Fix type signature for __divmod__

shoyer committed 361 days ago

7fb68cac

[shape_poly] Improve threefry with symbolic shapes

gnecula committed 361 days ago

bc3306c8

Merge pull request #25728 from zhenying-liu:scipy.misc

Google-ML-Automation committed 361 days ago

7997f080

Make inspect_array_sharding work without mesh context manager too.

yashk2810 committed 361 days ago

23eaf216

Merge pull request #25740 from jakevdp:remove-array-api

Google-ML-Automation committed 361 days ago

b304b9ef

Remove deprecated jax.experimental.array_api

jakevdp committed 361 days ago

c7b0d681

[Mosaic TPU] Add some elementwise canonicalizations

sharadmv committed 362 days ago

4caa263a

Remove the need for check_rep for with_sharding_constraint.

pschuh committed 362 days ago

b49ba655

fix the doc error: module 'scipy.misc' has no attribute 'face'

zhenying-liu committed 362 days ago

77c6947a

Rename pybind_extension to nanobind_extension.

hawkinsp committed 362 days ago

90d8f378

Suppress MSAN warnings from SVD that are showing up in CI.

hawkinsp committed 362 days ago

61dd0412

Merge pull request #25214 from jakevdp:einsum-optimize

Google-ML-Automation committed 362 days ago

52cc5c7f

Newer Older