ROCm/jax

Pull Requests Commits

The reshape call has a numerical bug where the target reshape size is always higher than array size if num_devices is not a nonnegative power of two. To fix this, this commit adds padding

gulsumgudukbay committed 106 days ago

1598f8ce

Merge pull request #31696 from mattjj:invalidate-debug-info-by-default

Google-ML-Automation committed 106 days ago

d94c8ac7

[jax] Deprecate `jax_pmap_no_rank_reduction` config state.

danielsuo committed 106 days ago

b26ade28

Merge pull request #31375 from jaro-sevcik:use-jacobi-for-lax-eigh

Google-ML-Automation committed 106 days ago

ee59f81d

Merge pull request #31707 from jburnim:jburnim_pallas_interpret_oob

Google-ML-Automation committed 106 days ago

0233945a

[IFRT] Define `UserContext::Id()`

hyeontaek committed 106 days ago

e4b8959e

Preferentially map the innermost physical axis for higher bandwidth on TPU7x.

Google-ML-Automation committed 106 days ago

182b4672

[Pallas] Improve error msg for out-of-bounds reads of blocks of inputs.

jburnim committed 106 days ago

5763a4fc

[debug-info] invalidate debug_info by default in jaxpr.replace

mattjj committed 106 days ago

eaef3b2d

Merge pull request #31685 from jax-ml:cond-lojax

Google-ML-Automation committed 106 days ago

dc635a6d

Add a lower-to-lojax rule for cond

dougalm committed 106 days ago

616993d3

[Pallas:MGPU][NFC] Refactor `scratch_view` to operate on a single struct.

allanrenucci committed 106 days ago

b904ae0a

Merge pull request #31695 from mattjj:is-high-avals

Google-ML-Automation committed 106 days ago

d870b546

[Mosaic GPU] Relax verification logic for TMEM refs.

allanrenucci committed 106 days ago

7a33a8b4

[hijax] make hi_prim.is_high(*avals, **params) take avals as args

mattjj committed 106 days ago

cf2321f9

Merge pull request #31694 from mattjj:lint-fix

Google-ML-Automation committed 106 days ago

1077e837

mattjj committed 106 days ago

c0ace7fa

Update XLA dependency to use revision http://github.com/openxla/xla/commit/0e3619a84b24405e5f717b09181a818016754ebf

Google-ML-Automation committed 107 days ago

7eb58950

PR #31653: [hijax] add hijax to stages

mattjj committed 107 days ago

d9ee7d86

Merge pull request #31686 from mattjj:stages-remove-args-flat

Google-ML-Automation committed 107 days ago

2426c4eb

mattjj committed 107 days ago

4328f731

Add an enable_x64 option to jaxlib batched_device_put.

hawkinsp committed 107 days ago

f3d5d1f3

Merge the two tests named jax_jit_test.py.

hawkinsp committed 107 days ago

25a9f289

Merge pull request #31682 from mattjj:hijax-box-constructor

Google-ML-Automation committed 107 days ago

c5df168b

Remove special handling of default_device from jax_jit.

hawkinsp committed 107 days ago

218cf308

[jax2tf] Removed an unused call to `convert` in tests

superbobry committed 107 days ago

a312407f

[hijax] make Box constructor bind new_box_p

mattjj committed 107 days ago

f1642e06

Remove jax_jit extra_jit_context.

hawkinsp committed 107 days ago

f9183dac

Remove jax_jit._is_float0.

hawkinsp committed 107 days ago

bf75b5ee

Merge pull request #31647 from jakevdp:lax-int4

Google-ML-Automation committed 107 days ago

1b45c232

Older