ROCm/jax

Pull Requests Commits

[Mosaic GPU][NFC] Implement `inline_mgpu` and `custom_primitive` for WG semantics.

dimitar-asenov committed 216 days ago

295b92b9

Add visibility to `//jax:pretty_printer`

Jake VanderPlas committed 216 days ago

603d27ea

Update XLA dependency to use revision

Google-ML-Automation committed 216 days ago

c8e2fb20

Fix jax caching bug. The memory space selection logic should take into account the memory space of the requested sharding.

pschuh committed 216 days ago

b6e71051

Merge pull request #30368 from froystig:layout-doc

Google-ML-Automation committed 216 days ago

d16bf37c

write docs about controlling array layouts

froystig committed 217 days ago

641a768b

Remove `SafeNumDevices` now that all shardings are implemented in C++

yashk2810 committed 217 days ago

3eaa7f90

Loosen test tolerance for `tests/linalg_sharding_test.py::LinalgShardingTest::test_batch_axis_sharding_jvp`.

ZixuanJiang committed 217 days ago

2616b4d3

Merge pull request #29248 from DanisNone:ufuncs

Google-ML-Automation committed 217 days ago

36ae9d1b

Fix mesh_cast + new rng keys usage

yashk2810 committed 217 days ago

5371d6b7

[Pallas] Only run TPU interpret mode tests on CPU.

jburnim committed 217 days ago

8e21ca66

If the input aval mesh is empty, use the mesh of out_sharding passed to with_sharding_constraint on the output aval. We only use the mesh and not the pspec. This leads to more tracing cache hits as demonstrated in the test.

yashk2810 committed 217 days ago

bd32431a

Fix `py_import` dependencies.

Google-ML-Automation committed 217 days ago

433a433a

[Pallas][Mosaic GPU] Fix racy test by having each warpgroup work on non-overlapping data.

Rifur13 committed 217 days ago

0562069c

[JAX] Purge more caches when clearing backends

hyeontaek committed 217 days ago

ffd90dae

Migrate third party uses of `jax.lib.xla_bridge.get_backend` to `jax.extend.backend.get_backend`.

danielsuo committed 217 days ago

858c6f65

[Pallas] Make semantics of BlockSpec.memory_space in emit_pipeline consistent with pallas_call.

justinjfu committed 217 days ago

6057a1d9

Merge pull request #30356 from jakevdp:complex-fftfreq

Google-ML-Automation committed 217 days ago

59f86caf

[jex] Add `proto` <-> `xla_client.HloSharding` and `proto -> `xla_client.OpSharding` APIs to jex.

danielsuo committed 217 days ago

e577ef39

Fetch XLA archive using the GitHub API endpoint instead of the web link

nitins17 committed 217 days ago

6fee019a

Merge pull request #28968 from olupton:cuda-13-build-system

Google-ML-Automation committed 217 days ago

0b2caf9b

Allow setting transfer_size for cross host transfers.

pschuh committed 217 days ago

81317ff5

Merge pull request #29343 from SuriyaaMM:gumbel_impl

Google-ML-Automation committed 217 days ago

ccafdaa2

jnp.fft.fftfreq: fix regression in complex dtype support

jakevdp committed 217 days ago

f26c2cab

Automated g4 rollback of changelist 784947368.

Google-ML-Automation committed 217 days ago

58fb923f

Consider this example:

yashk2810 committed 217 days ago

b97ad423

[Mosaic] Add store canonicalization for an expand reshape->store fusion

voznesenskym committed 217 days ago

7cdc3671

gumbel distribution implementation

SuriyaaMM committed 217 days ago

c20e1ee2

Merge pull request #30351 from jax-ml:direct-linearize-changelog-notes

Google-ML-Automation committed 217 days ago

bfa581b7

Add note about direct linearize to changelog

dougalm committed 217 days ago

51ce5c15

Newer Older