ROCm/jax

Pull Requests Commits

Test - Memory use

pelumi1163 committed 18 days ago

6b4254fd

Test - Memory use

pelumi1163 committed 18 days ago

5c058fe3

Test - Memory use

pelumi1163 committed 18 days ago

c80a42c1

Test - Memory use

pelumi1163 committed 18 days ago

ecd78036

Test - Memory use

pelumi1163 committed 18 days ago

ceb30c19

Test - Memory use

pelumi1163 committed 18 days ago

ec34fa5f

Memory use check

pelumi1163 committed 18 days ago

a5e292e2

add chmod for /usr/local

pelumi1163 committed 22 days ago

5d039a9b

Remove docker process

pelumi1163 committed 22 days ago

56981218

Test Static CPU pods

pelumi1163 committed 23 days ago

59fe0336

Test Static CPU pods

pelumi1163 committed 23 days ago

983f7cbd

Test Static CPU pods

pelumi1163 committed 23 days ago

7af2890d

Merge branch 'jax-ml:main' into amd-main

magaonka-amd committed 25 days ago

Verified a86c0159

Roll back #39282 due to CI failures

nitins17 committed 25 days ago

29fd8792

Merge pull request #39282 from Intel-tensorflow:mraunak/presubmit_oneapi

Google-ML-Automation committed 25 days ago

b00b9784

Merge pull request #39450 from kodlan:fix/issue-7465-vmap-kwargs-in-axes

Google-ML-Automation committed 25 days ago

f3a93cbd

[pallas] The wrapper function in `pl.with_scoped` now accepts `kwargs`

superbobry committed 25 days ago

3a055174

Merge pull request #39470 from mattjj:push-rwxyltxntpkk

Google-ML-Automation committed 25 days ago

34cd9238

Merge pull request #39353 from Gattocrucco:faster-poisson-sampler

Google-ML-Automation committed 25 days ago

33b1f842

[custom_vjp3] fix custom_jvp + pickle interaction

mattjj committed 25 days ago

4fa61baa

Merge pull request #39467 from mattjj:push-xkrtkyustwzy

Google-ML-Automation committed 25 days ago

7238390c

[docs] flip the order of vjp3 advanced subsections

mattjj committed 25 days ago

ddf29ce1

Mention keyword arguments in the vmap in_axes arity error

kodlan committed 25 days ago

e702c77f

Expose SparseCore SMEM

vsytch committed 25 days ago

a9dc4c84

Prevent the TPU backend init warning from surfacing as an error in random pytest tests.

belitskiy committed 25 days ago

3b090a42

Merge pull request #39461 from jakevdp:dslice-test

Google-ML-Automation committed 25 days ago

7451f711

Merge pull request #39447 from mattjj:push-knkykrxvluow

Google-ML-Automation committed 25 days ago

84f4a21d

Merge pull request #39462 from hawkinsp:mpmath

Google-ML-Automation committed 25 days ago

19981a7f

[CI] Shard custom_api_test on GPUs to fix the test timeout.

belitskiy committed 25 days ago

b80e7310

Add validation to SDS constructor to error out if shape has None in it.

yashk2810 committed 25 days ago

01ce5619

Older