ROCm/jax

Pull Requests Commits

[ROCm] Stage libhipsparse in the ROCm wheel test dependencies

magaonka-amd committed 39 days ago

1cc45223

Reverts 1b61209a718fa6e7da6f267361ce3ed118d52fc6

Google-ML-Automation committed 39 days ago

aec6fc5e

Check memory_space constraint even if no mesh exists (cause it has nothing to do with meshes)

yashk2810 committed 39 days ago

8736c2b1

[Pallas:MGPU] Fix signature of print_layout methods

apaszke committed 39 days ago

d37f6c8d

[Pallas:MGPU] Erase swizzle from transform list before handling transforms

apaszke committed 39 days ago

be0dd453

[Mosaic GPU] Take bitwidth into account when inferring TMEM load layout

apaszke committed 39 days ago

f857b4c6

Reverts ee61766245915256061f5aca6cd342ff1b4957b1

superbobry committed 39 days ago

1b61209a

[Mosaic GPU] Add transfer plan analysis for 8-bit transposed txmatrix

apaszke committed 39 days ago

ebda9f6d

Update XLA dependency to use revision http://github.com/openxla/xla/commit/0e45fa4e6960ed86aceb51c0aec071e299a6cc31

Google-ML-Automation committed 39 days ago

a74b888b

[pallas:mosaic] Yet another fix for named grid and batching

superbobry committed 39 days ago

7a46165b

Migrate JAX-Triton custom calls to internal XLA FFI headers

EusebioDM committed 39 days ago

c6ab31b9

Exclude drop_fields fields when auto-populating data/meta fields.

marksandler2 committed 39 days ago

eddef19c

Only check `.mat` equivalency when the aval is a ShapedArray in scan abstract eval. Fixes https://github.com/jax-ml/jax/issues/39700

yashk2810 committed 39 days ago

75a26ecf

Reverts 3ba417533f96ee1a6022b93d173456efde0e253b

Google-ML-Automation committed 39 days ago

f7fef7a0

Merge pull request #39697 from mattjj:push-owyzwvzyruwx

Google-ML-Automation committed 39 days ago

7712ef95

add test for eval_jaxpr_p partial eval rule

mattjj committed 39 days ago

a712434f

Delete compute_on2 now that everyone is on compute_on

yashk2810 committed 39 days ago

1558b651

Expose `compute_on` as an alias of `compute_on2` in the preparation of deleting compute_on2 as the name.

yashk2810 committed 40 days ago

ad1df344

[Pallas TPU] Expose accumulator references to be used with explicit MXU APIs.

bchetioui committed 40 days ago

3ba41753

Improve "unsupported reshape merge" error message in fuser

sharadmv committed 40 days ago

44d7bc0a

Reverts d67bc6927d29f713f19d40df72fd1b7a3ff88a8e

Google-ML-Automation committed 40 days ago

ee617662

[Pallas] Expand the error message in `get_tpu_info` to hint towards defining an `AbstractMesh` when _not_ running on a TPU device.

bchetioui committed 40 days ago

1ebd186c

Merge pull request #39685 from jakevdp:fix-ratchet

Google-ML-Automation committed 40 days ago

43dd7542

Merge pull request #39655 from ashahba:fix/oneapi-workflow-permissions-it

Google-ML-Automation committed 40 days ago

965efd9f

Remove an extra weakref_lru_cache added by AI

yashk2810 committed 40 days ago

7f4e0282

[actions] remove ratchet versions & pin specific versions instead

jakevdp committed 40 days ago

54b1be2b

[NFC] Refactor layout transforms

apaszke committed 40 days ago

68970bf5

[oneAPI] Fix startup_failure by removing permissions escalation in called workflow

ashahba committed 40 days ago

523c4757

Merge pull request #39591 from jax-ml:dependabot/github_actions/aws-actions/configure-aws-credentials-6.2.3

Google-ML-Automation committed 40 days ago

73befe48

Merge pull request #39592 from jax-ml:dependabot/github_actions/actions/setup-python-7.0.0

Google-ML-Automation committed 40 days ago

a13a4a7f

Older