ROCm/jax

Pull Requests Commits

Use gha artifacts to pass the rocm plugins through for pytest job

alekstheod committed 79 days ago

cd9a6d7c

[Mosaic GPU] Add bitwidth to IsTransferable.

golechwierowicz committed 79 days ago

b94ce497

[Pallas:MGPU] Expose the leader_tracked in copy_gmem_to_smem

apaszke committed 79 days ago

d8bbe601

Merge pull request #35796 from gnecula:ser_abs_dev2

Google-ML-Automation committed 79 days ago

422cf9d2

[export] Add support for serializating AbstractMesh.abstract_device

gnecula committed 80 days ago

e812c59c

[Pallas] Add excluding_{lhs,rhs}_dims to jax_dot_dims_to_tpu_dot_dot_dims.

WindQAQ committed 80 days ago

2acb4492

Merge pull request #35793 from mattjj:is-vjp

Google-ML-Automation committed 80 days ago

d7ccdd1e

[JAX] Require all positional arguments passed to bind to be JAX values.

hawkinsp committed 80 days ago

c43f0a0f

extremely complex fix

mattjj committed 80 days ago

04581201

jaxlib: Also let the message attribute be optional for JaxRuntimeError.

pizzud committed 80 days ago

284116c7

Remove references to has_cpu_sparse, which is always true.

hawkinsp committed 80 days ago

5a9490fa

[Mosaic] Generate enums in Python bindings

tlongeri committed 80 days ago

ebf39ff6

jaxlib: Expose the status attributes on XlaRuntimeError.

pizzud committed 80 days ago

ffcc5ea9

[jaxlib] Fixed usages of `typing_extensions.CapsuleType` in stubs

superbobry committed 80 days ago

ed1e8d37

Merge pull request #35713 from alekstheod:patch-2

Google-ML-Automation committed 80 days ago

44e12c3a

Merge pull request #35747 from jakevdp:ai-policy

Google-ML-Automation committed 80 days ago

719aaea1

[jaxlib] Added a missing include and regenerated the .pyis

superbobry committed 80 days ago

062845f0

Increase size of gpu_pallas_interpret_test_cpu to avoid timeout.

mwhittaker committed 80 days ago

3853b241

Expose register_extra_dialect in mosaic to support loading custom dialect attributes for backward compatability.

Google-ML-Automation committed 80 days ago

78126876

Increase size of `lax_scipy_spectral_dac_test` to avoid timeouts.

mwhittaker committed 80 days ago

3f6a1731

Merge pull request #35748 from jakevdp:pyrefly-sparse

Google-ML-Automation committed 80 days ago

da78e804

[MGPU] Add AsyncStoreSmemToTmem Op.

golechwierowicz committed 80 days ago

336a1de7

[Mosaic GPU] Add support for reducing strided and splat reduce layout constaints.

allanrenucci committed 80 days ago

18a19f8b

[Pallas/Mosaic GPU] Enable warpgroup semantics-based tests for the `HelpersTest` and `DistributedTest` classes.

bchetioui committed 80 days ago

0453525b

[Pallas] Add more test cases for folding last two dims reshape.

WindQAQ committed 80 days ago

2b93660d

[Pallas/Mosaic GPU] Enable lowering `semaphore_{read,signal,wait}_p` using warpgroup semantics.

bchetioui committed 80 days ago

4e9175fb

Remove extra space in JAX OSS workflow job names.

mwhittaker committed 80 days ago

cc60de73

Merge pull request #35710 from alekstheod:patch-1

Google-ML-Automation committed 80 days ago

212235b3

[pyrefly] add type checking for jax/experimental/sparse

jakevdp committed 80 days ago

f722be11

[Pallas:TPU] Use constant 0 indices for grid dims known to be 1

apaszke committed 80 days ago

7da7f39c

Older