ROCm/jax

Pull Requests Commits

Fix deadlock when computing cached Sharding::type() values.

hawkinsp committed 262 days ago

8a6efa31

Automated Code Change

Google-ML-Automation committed 262 days ago

7874d79f

Automated Code Change

Google-ML-Automation committed 262 days ago

d1c7ba43

Automated Code Change

Google-ML-Automation committed 262 days ago

3f083cae

Automated Code Change

Google-ML-Automation committed 262 days ago

477b1082

Automated Code Change

Google-ML-Automation committed 262 days ago

ad36f7f2

Update XLA dependency to use revision

Google-ML-Automation committed 263 days ago

b1b54f9b

Update XLA dependency to use revision

Google-ML-Automation committed 264 days ago

2e62693f

Check that memory_kind of an aval is always None

yashk2810 committed 264 days ago

fc5d9a4f

[pallas:mosaic_gpu] Fixed a typo in `_barrier_arrive_pp_eqn`

superbobry committed 264 days ago

aab66139

[Mosaic GPU] Limit the maximum number of registers per thread to 255.

Rifur13 committed 264 days ago

d81c0ffe

Don't set `memory_kind` to `None` if the mesh is AbstractMesh and the

yashk2810 committed 264 days ago

549f1cd8

PR #26906: [jax.distributed] Allow explicitly setting slice_index

gspschmid committed 264 days ago

5d4ac775

Fix a possible race in pjit.cc.

hawkinsp committed 264 days ago

be1a554d

Merge pull request #27731 from Qazalbash:fix-doc-grad-checkpoint

Google-ML-Automation committed 264 days ago

027f7333

Parameterize the random tests taking out_sharding argument in pjit_test.py

yashk2810 committed 264 days ago

e6b01bd1

Always force synchronous pipelining when we have vmem storage and trivial

Google-ML-Automation committed 264 days ago

e2f67e0e

Merge pull request #27730 from froystig:out-shard-trunc-normal

Google-ML-Automation committed 264 days ago

5a7dc42a

`_attempt_rewriting_take_via_slice()`: canonicalize the slice index before checking it's not too long, so that e.g. `my_1d_array[:, ...]` can be treated as a slice rather than generating a gather operation.

Google-ML-Automation committed 264 days ago

35d75183

[mgpu:pallas] Fix swizzling check bug where it was comparing w/ #bytes rather than #elems.

cperivol committed 264 days ago

b9007145

[mgpu] Foreach to handle scalar registers in fragmented arrays.

cperivol committed 264 days ago

53abbd56

[mgpu:pallas] Swizzle elements computed using bitwidth rather than bytewidth.

cperivol committed 264 days ago

da7b1577

Update XLA dependency to use revision

Google-ML-Automation committed 265 days ago

5a29311c

[mgpu:pallas] Typo in `UnswizzleRef.untransform_reshape()` check.

cperivol committed 265 days ago

cbae2539

[pallas:mgpu] Check that swizzle dim is not transposed in copy_smem_to_gmem()

cperivol committed 265 days ago

e4a381c1

[Mosaic GPU] Allow replicating data over warps

apaszke committed 265 days ago

635805e9

Merge pull request #27707 from vfdev-5:fix-named-sharding-deadlock

Google-ML-Automation committed 265 days ago

bb638506

[Mosaic GPU] Don't force TiledLayout.lane_dims to partition data

apaszke committed 265 days ago

b0a920dd

[pallas:mosaic_gpu] Added pretty printing to primitives consuming refs

superbobry committed 265 days ago

206dec85

Fixed deadlock in NamedSharding ctor

vfdev-5 committed 265 days ago

1b63d5e2

Older