ROCm/jax

Pull Requests Commits

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

danielsuo committed 166 days ago

858c6f65

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

justinjfu committed 166 days ago

6057a1d9

Merge pull request #30356 from jakevdp:complex-fftfreq

Google-ML-Automation committed 166 days ago

59f86caf

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

danielsuo committed 166 days ago

e577ef39

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

nitins17 committed 166 days ago

6fee019a

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

Google-ML-Automation committed 166 days ago

0b2caf9b

Allow setting transfer_size for cross host transfers.

pschuh committed 166 days ago

81317ff5

Merge pull request #29343 from SuriyaaMM:gumbel_impl

Google-ML-Automation committed 166 days ago

ccafdaa2

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

jakevdp committed 166 days ago

f26c2cab

Automated g4 rollback of changelist 784947368.

Google-ML-Automation committed 166 days ago

58fb923f

Consider this example:

yashk2810 committed 166 days ago

b97ad423

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

voznesenskym committed 166 days ago

7cdc3671

gumbel distribution implementation

SuriyaaMM committed 166 days ago

c20e1ee2

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

Google-ML-Automation committed 166 days ago

bfa581b7

Add note about direct linearize to changelog

dougalm committed 166 days ago

51ce5c15

Merge pull request #30348 from NeilGirdhar:update_line_length

Google-ML-Automation committed 166 days ago

4a0a500c

Merge pull request #29561 from NeilGirdhar:fix_initializers

Google-ML-Automation committed 166 days ago

9e57de84

Update Ruff line length and python version.

NeilGirdhar committed 166 days ago

Verified 3e3a31e0

Merge pull request #30248 from olupton:multi-gpu-test-on-single-gpu-machine

Google-ML-Automation committed 166 days ago

c331e7bb

Polish initializer annotations

NeilGirdhar committed 166 days ago

Verified 8edec0ea

#sdy Add section on Shardy migration issues about all meshes being the same and custom_partitioning.

bartchr808 committed 166 days ago

78dd54d3

update `use_shardy_partitioner` description.

tomnatan30 committed 166 days ago

ecad5a88

improve AOT input sharding/layout mismatch error message

froystig committed 167 days ago

7c885373

Update XLA dependency to use revision

Google-ML-Automation committed 167 days ago

7414f8ce

Fix creation of sharding rules for fused_attention_stablehlo.

tomnatan30 committed 167 days ago

b1af72db

[JAX] Enable Shardy by default in JAX.

tomnatan30 committed 167 days ago

df8e17a5

Update XLA dependency to use revision

Google-ML-Automation committed 168 days ago

109bcae8

Merge pull request #29370 from mattjj:mutable-array-custom-vjp-scan3

Google-ML-Automation committed 168 days ago

860a29d7

Merge pull request #30340 from mattjj:mutable-array-shmap

Google-ML-Automation committed 168 days ago

213fa6fc

Separate out `vma` and `unreduced/reduced` in the ShapedArray str_short().

yashk2810 committed 168 days ago

0acbd3b3

Newer Older