ROCm/jax

Pull Requests Commits

Add AMD ROCm GPU CI post-build check (#137)

charleshofer committed 320 days ago

ebf4a54f

bazel visibility change

Jake VanderPlas committed 349 days ago

640cb009

Make sharding_in_types work with Shardy

yashk2810 committed 349 days ago

b2b38679

Respect the original memory kind on reshape, transpose and replicate methods of PositionalSharding. Fixes https://github.com/jax-ml/jax/issues/25769

yashk2810 committed 349 days ago

fb832afc

#sdy add repr for Sdy ArraySharding and DimSharding

bartchr808 committed 349 days ago

cbcc883e

Merge pull request #25786 from vfdev-5:add-313-ft-configs

Google-ML-Automation committed 349 days ago

196eec82

Make api_test.py work when test cases are run using multiple threads.

hawkinsp committed 349 days ago

e20523c2

Added 3.13 ft requirements lock file and updated WORKSPACE

vfdev-5 committed 349 days ago

00806dda

Implement the extension to the custom_partitioning API.

bixia1 committed 349 days ago

c4ac0dd6

Match the behavior on single host wrt multi-host if tiled=False. Fixes https://github.com/jax-ml/jax/issues/25783

yashk2810 committed 349 days ago

b3833dc7

Merge pull request #25527 from vfdev-5:single-python-version-build-py

Google-ML-Automation committed 349 days ago

6e1f060a

[Mosaic GPU] Allow multiple indexing on refs

justinjfu committed 349 days ago

d99a637d

[sharding_in_types] Functions like einsum, reshape, broadcast_in_dim, broadcasted_iota, convert_element_type and sharding_cast that take out_sharding as an argument in their signature should also allow `PartitionSpec` instead of just `NamedSharding` as an input.

yashk2810 committed 349 days ago

3848f0d2

[Pallas] Add empty/empty_like helper functions

sharadmv committed 349 days ago

c1a60c67

Update XLA dependency to use revision

Google-ML-Automation committed 349 days ago

55119490

#sdy Move Shardy mesh lift inlining pass after verification.

bartchr808 committed 349 days ago

5c097c8f

Add a unittest test extension that runs test cases in parallel using threads.

hawkinsp committed 349 days ago

0389d617

Port tests away from setUpClass and setUpModule to setUp alone.

hawkinsp committed 349 days ago

3fa55728

[pallas:mosaic_gpu] Fix the tests following the changes to `pl.core_map`

superbobry committed 349 days ago

f1f98afe

[JAX] Add a new jax_num_cpu_devices flag that allows the user to specify the number of CPU directly.

hawkinsp committed 349 days ago

51b9fe30

[Mosaic TPU] Be much more aggressive in inferring large 2nd minor layouts for 16-bit types on v6

apaszke committed 349 days ago

f96339be

[Mosaic TPU] Add support for second minor broadcasts with packed types

apaszke committed 349 days ago

5fd1b2f8

[Mosaic TPU] Add support for true divide in bf16 on TPUv6

apaszke committed 349 days ago

e954930e

[Mosaic] Use tpu::CreateMask for getX32VmaskByPaddingEnd.

WindQAQ committed 349 days ago

bf94389b

Merge pull request #25754 from andportnoy:patch-4

Google-ML-Automation committed 349 days ago

4718121e

Removed leftover mentions of xmap from the code

superbobry committed 349 days ago

90201ce2

Merge pull request #25594 from zhenying-liu:activation-offloading-doc

Google-ML-Automation committed 349 days ago

81db3219

Add JAX events that have time spans, not only durations.

Google-ML-Automation committed 349 days ago

1bd781d9

Merge branch 'jax-ml:main' into activation-offloading-doc

zhenying-liu committed 350 days ago

Verified 21fb171e

Merge pull request #25761 from jakevdp:array-api-update

Google-ML-Automation committed 350 days ago

6d08f36f

Older