ROCm/jax

Pull Requests Commits

Merge pull request #25569 from hawkinsp:numpyver

Google-ML-Automation committed 1 year ago

464e5a27

More linearize fixes

dougalm committed 1 year ago

2be9a699

Migrate mhlo dialect extension to nanobind.

hawkinsp committed 1 year ago

3f24dfd2

Partial discharge for scan_p ops.

cperivol committed 1 year ago

aaabb975

Use capture_stderr instead of packing sys.stderr

apaszke committed 1 year ago

2259a13d

[Mosaic TPU] Add support for bf16 second minor reductions in TPUv6

apaszke committed 1 year ago

6edfe9ea

Enable --config=clang only on newer Clang versions

nitins17 committed 1 year ago

bcca77cd

[Mosaic GPU] Clean up imports in `gpu_dialect_test.py`.

bchetioui committed 1 year ago

6a03ea3e

Migrate StableHLO Python extension to nanobind.

hawkinsp committed 1 year ago

3d54d035

Merge pull request #25412 from shangz-ai:patch-3

Google-ML-Automation committed 1 year ago

66c8a2e9

Use the right health check suppression

apaszke committed 1 year ago

3c52ca6e

[pallas:mosaic_gpu] Added a lowering rule for the general `lax.while_loop_p`

superbobry committed 1 year ago

5e9cfce2

[mosaic_gpu] Error on static out of bounds indices in `utils.parse_indices`

superbobry committed 1 year ago

98067fc1

Merge pull request #23830 from pearu:pearu/acos

Google-ML-Automation committed 1 year ago

5beb4794

Increase the timeout in profiler tests, but use threading.Event to exit as soon as possible

apaszke committed 1 year ago

7c1bea6b

[Mosaic TPU] Add support for exp, exp2 and log in bf16 on TPUv6

apaszke committed 1 year ago

d95b95b4

Merge pull request #25559 from hawkinsp:minver

Google-ML-Automation committed 1 year ago

5c9756bb

Use StableHLO acos and update complex acos accuracy tests.

pearu committed 1 year ago

f592173c

Increase the minimum NumPy version to v1.25.

hawkinsp committed 1 year ago

ee457184

Bump xla_extension_version after jaxlib release.

hawkinsp committed 1 year ago

7776982a

Reverts b56dc63160eaccd7df05d03b1c38f804ff85f564

Google-ML-Automation committed 1 year ago

25524abc

Merge pull request #25449 from mattjj:ref-errors-3

Google-ML-Automation committed 1 year ago

96d4a758

[Mosaic:TPU] Allow null parts for tpu.pack_subelements, meaning "don't care"

tlongeri committed 1 year ago

dc0b7747

mattjj committed 1 year ago

42ac4ca3

Merge pull request #24516 from mattjj:improve-concreteness-error-in-remat-3

Google-ML-Automation committed 1 year ago

32627702

[Mosaic:TPU] Fix bug after cl/707025084

tlongeri committed 1 year ago

f9737b95

improve checkpoint / remat concreteness error with static_argnums

mattjj committed 1 year ago

9acd4a95

[sharding_in_types] Add tests allowing inputs and outputs of jit to have different axis_types on their mesh than the axis_types on the surrounding mesh context

yashk2810 committed 1 year ago

09fdd0da

Allow `P.UNCONSTRAINED` in out_shardings at top level jit. This is required for sharding in types to work properly when out_avals contain UNCONSTRAINED specs.

yashk2810 committed 1 year ago

e854f165

Merge pull request #25555 from jakevdp:combine-avals

Google-ML-Automation committed 1 year ago

b56dc631

Newer Older