ROCm/jax

Pull Requests Commits

Merge pull request #31328 from olupton:test-dump-on-deserialize

Google-ML-Automation committed 30 days ago

174d514d

[Mosaic GPU][NFC] Add a test for legacy Mosaic GPU custom calls.

allanrenucci committed 30 days ago

eae3b49b

[Mosaic GPU] Reduce test size to make it run on GPUs with less SMEM (e.g. RTX Pro 6000 Blackwell)

dimitar-asenov committed 30 days ago

4fe443e3

Merge pull request #33724 from jakevdp:dtype-dep

Google-ML-Automation committed 30 days ago

f68ae2b3

Merge pull request #33608 from partev:patch-2

Google-ML-Automation committed 30 days ago

b116a4a0

Change //jax:experimental to be a target that reexports other targets.

hawkinsp committed 30 days ago

53a2a5c8

Fix deprecated use of np.dtype

jakevdp committed 30 days ago

b82aa387

Add jax_check_static_indices configuration

jakevdp committed 30 days ago

6832479f

[pallas:mosaic] Use `tree_util.register_dataclass` for `BufferedRef`

superbobry committed 30 days ago

9b31b8ef

Test dumping HLO when reading compilation cache

olupton committed 30 days ago

1588d16e

Merge pull request #33710 from jax-ml:remove-tracing-wrappers

Google-ML-Automation committed 30 days ago

92a5b3bc

Remove //jax:extend bazel target (use //jax/extend/... instead).

hawkinsp committed 30 days ago

03132bdc

Fixes for Jetson Thor devices

olupton committed 30 days ago

01513b27

[Mosaic GPU] Fix broken tests when jax is coupled with an old jax lib version.

dimitar-asenov committed 30 days ago

d0fd8f08

Fix paths triggering H100/B200 presubmits.

bchetioui committed 30 days ago

fdfffdc3

[Mosaic GPU][NFC] Clarify comments on custom call variants.

allanrenucci committed 30 days ago

ca87f344

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

Google-ML-Automation committed 30 days ago

7f343072

Fix Python path for executing "hi" JAX code.

chr1sj0nes committed 30 days ago

c5f92284

Change int2 shape in pallas test

hanrach9 committed 30 days ago

80967f3d

Update PRNG documentation with corrected links

partev committed 30 days ago

Verified daedbcc0

[Pallas] Make rng tests inherit from JaxTestCase

justinjfu committed 30 days ago

10a43df9

Merge pull request #33706 from mattjj:more-code-than-dougal2

Google-ML-Automation committed 30 days ago

fb0b1c0e

Remove lax_metal_test and associated CI job

jakevdp committed 30 days ago

aa886ea9

[no-thunks] Deprecate old form of `cond` and remove some tracing layers.

dougalm committed 30 days ago

4cbcbe7e

add ensure_compile_time_eval support to custom_jvp/vjp

mattjj committed 30 days ago

7bda7a05

fix printing of unreduced Arrays

mattjj committed 30 days ago

ec135957

Remove CUDA12 tests with jaxlib built from HEAD from the mandatory presubmit.

ybaturina committed 30 days ago

245b2c54

Reverts 23cd412cebf6ac17b27b3d0699aced5c9002ae89

dougalm committed 30 days ago

19ab33dd

hawkinsp committed 30 days ago

92c283e9

Merge pull request #33702 from mattjj:more-code-than-dougal

Google-ML-Automation committed 30 days ago

08e75245

Newer Older