ROCm/jax

Pull Requests Commits

Merge pull request #27198 from jakevdp:lax-docs

Google-ML-Automation committed 278 days ago

30941480

Make blocked_fold_in consistent when the block sizes induce padding

Google-ML-Automation committed 278 days ago

13541e9f

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

charleshofer committed 278 days ago

Verified c46b4fc0

jax.lax: improve docs for several APIs

jakevdp committed 278 days ago

8b46e53a

Added exit 1 if git patch is failed + other checks

vfdev-5 committed 278 days ago

9145d617

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

charleshofer committed 278 days ago

c7b407c9

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

bchetioui committed 278 days ago

1e36cbe5

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

apaszke committed 278 days ago

8da93249

Update XLA dependency to use revision

Google-ML-Automation committed 278 days ago

7a459f0e

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

bchetioui committed 278 days ago

ba2f7c9a

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

apaszke committed 278 days ago

d4bd2570

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

chr1sj0nes committed 278 days ago

38d52a19

[Mosaic GPU] Remove sub-byte conversion restriction

apaszke committed 278 days ago

34cd5b0d

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

yashk2810 committed 278 days ago

549973de

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

changhuilin committed 278 days ago

f174b00f

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

emilyfertig committed 278 days ago

8c351917

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

superbobry committed 278 days ago

051687dc

Compute tile index using tile-based coordinates

Google-ML-Automation committed 278 days ago

b4966130

Merge pull request #27164 from MichaelHudgins:a4-testing

Google-ML-Automation committed 278 days ago

b74b16f9

Add B200 testing to continuous workflow

MichaelHudgins committed 278 days ago

ecf7fde7

Fix error in pallas tutorial

Google-ML-Automation committed 278 days ago

4f704713

Replace cached function get_replicated_hlo_sharding() with a constant.

hawkinsp committed 278 days ago

20658fab

Merge pull request #26980 from carlosgmartin:categorical_replace

Google-ML-Automation committed 278 days ago

ebcae0d3

Remove code that preserved _original_py_fns on C++ classes.

hawkinsp committed 278 days ago

be5d13af

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

bchetioui committed 278 days ago

9a686e0b

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

carlosgmartin committed 278 days ago

3f59fa68

Fix auditwheel version issue (#288)

mrodden committed 279 days ago

Verified d864b4fb

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

Google-ML-Automation committed 279 days ago

de9ad6ba

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

nitins17 committed 279 days ago

031614c2

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

apaszke committed 279 days ago

3649da56

Newer Older