ROCm/jax

Pull Requests Commits

Enable upstream CI on release branches

charleshofer committed 1 year ago

c4f188cc

Merge pull request #294 from ROCm/ci-upstream-sync-151_1

charleshofer committed 1 year ago

Verified c46b4fc0

Merge branch 'rocm-main' into ci-upstream-sync-151_1

charleshofer committed 1 year ago

c7b407c9

[Mosaic GPU] Raise a `NotImplementedError` if `swizzle=16`.

bchetioui committed 1 year ago

1e36cbe5

[Mosaic GPU] Fuse slicing into s4 -> bf16 upcasts

apaszke committed 1 year ago

8da93249

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7a459f0e

[Mosaic GPU] Add transform inference rule for `mgpu.slice_smem`.

bchetioui committed 1 year ago

ba2f7c9a

[Mosaic GPU] Add a specialized layout for loading 4-bit inputs in WGMMA friendly layouts

apaszke committed 1 year ago

d4bd2570

[mosaic_gpu] Force flush all cupti activity, then unsubscribe.

chr1sj0nes committed 1 year ago

38d52a19

[Mosaic GPU] Remove sub-byte conversion restriction

apaszke committed 1 year ago

34cd5b0d

Allow pspec to be passed to device_put if there is a mesh in the surrounding context

yashk2810 committed 1 year ago

549973de

Replace the uses of `PjRtClient::Compile()` with `PjRtClient::CompileAndLoad()`.

changhuilin committed 1 year ago

f174b00f

Enable `jax.device_put` to a sharding with no local devices.

emilyfertig committed 1 year ago

8c351917

[pallas] `pallas_call_p` is now parameterized by a mesh

superbobry committed 1 year ago

051687dc

Compute tile index using tile-based coordinates

Google-ML-Automation committed 1 year ago

b4966130

Merge pull request #27164 from MichaelHudgins:a4-testing

Google-ML-Automation committed 1 year ago

b74b16f9

Add B200 testing to continuous workflow

MichaelHudgins committed 1 year ago

ecf7fde7

Fix error in pallas tutorial

Google-ML-Automation committed 1 year ago

4f704713

Replace cached function get_replicated_hlo_sharding() with a constant.

hawkinsp committed 1 year ago

20658fab

Merge pull request #26980 from carlosgmartin:categorical_replace

Google-ML-Automation committed 1 year ago

ebcae0d3

Remove code that preserved _original_py_fns on C++ classes.

hawkinsp committed 1 year ago

be5d13af

[Mosaic GPU] Add initial transform inference rules for `vector.{load,store}`.

bchetioui committed 1 year ago

9a686e0b

Add replace option to random.categorical to enable sampling without replacement.

carlosgmartin committed 1 year ago

3f59fa68

Fix auditwheel version issue (#288)

mrodden committed 1 year ago

Verified d864b4fb

Merge pull request #27157 from mar-muel:improve-random-choice-performance

Google-ML-Automation committed 1 year ago

de9ad6ba

Pin numpy~=2.1.0 in workflow file instead of test-requirements.txt

nitins17 committed 1 year ago

031614c2

[Mosaic GPU] Make the s4 -> bf16 upcast more flexible when it comes to vector length

apaszke committed 1 year ago

3649da56

Removed trivial docstrings from JAX tests

superbobry committed 1 year ago

0ff23404

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

55812c5d

[pallas:mosaic_gpu] `jnp.reduce_sum` now works for >1D arrays

superbobry committed 1 year ago

a7e5eaee

Older