ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

c8e2fb20

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

pschuh committed 1 year ago

b6e71051

Merge pull request #30368 from froystig:layout-doc

Google-ML-Automation committed 1 year ago

d16bf37c

write docs about controlling array layouts

froystig committed 1 year ago

641a768b

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

yashk2810 committed 1 year ago

3eaa7f90

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

ZixuanJiang committed 1 year ago

2616b4d3

Merge pull request #29248 from DanisNone:ufuncs

Google-ML-Automation committed 1 year ago

36ae9d1b

Fix mesh_cast + new rng keys usage

yashk2810 committed 1 year ago

5371d6b7

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

jburnim committed 1 year 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 1 year ago

bd32431a

Fix `py_import` dependencies.

Google-ML-Automation committed 1 year ago

433a433a

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

Rifur13 committed 1 year ago

0562069c

[JAX] Purge more caches when clearing backends

hyeontaek committed 1 year ago

ffd90dae

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

danielsuo committed 1 year ago

858c6f65

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

justinjfu committed 1 year ago

6057a1d9

Merge pull request #30356 from jakevdp:complex-fftfreq

Google-ML-Automation committed 1 year ago

59f86caf

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

danielsuo committed 1 year ago

e577ef39

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

nitins17 committed 1 year ago

6fee019a

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

Google-ML-Automation committed 1 year ago

0b2caf9b

Allow setting transfer_size for cross host transfers.

pschuh committed 1 year ago

81317ff5

Merge pull request #29343 from SuriyaaMM:gumbel_impl

Google-ML-Automation committed 1 year ago

ccafdaa2

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

jakevdp committed 1 year ago

f26c2cab

Automated g4 rollback of changelist 784947368.

Google-ML-Automation committed 1 year ago

58fb923f

Consider this example:

yashk2810 committed 1 year ago

b97ad423

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

voznesenskym committed 1 year ago

7cdc3671

gumbel distribution implementation

SuriyaaMM committed 1 year ago

c20e1ee2

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

Google-ML-Automation committed 1 year ago

bfa581b7

Add note about direct linearize to changelog

dougalm committed 1 year ago

51ce5c15

Merge pull request #30348 from NeilGirdhar:update_line_length

Google-ML-Automation committed 1 year ago

4a0a500c

Merge pull request #29561 from NeilGirdhar:fix_initializers

Google-ML-Automation committed 1 year ago

9e57de84

Older