ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

b372ce4b

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

4d8751bf

[Mosaic][Easy] - Wire up kernel names to MLIR dump

Google-ML-Automation committed 1 year ago

e53ff2cb

Fix member access to xla backend. The correct member is `client` instead of `backend`

Google-ML-Automation committed 1 year ago

b259fde5

When host local inputs on all hosts are the same, use `_DeferredShardArg` to do the transfers instead of `jit` to avoid blocking.

yashk2810 committed 1 year ago

8699f5d9

Merge pull request #25070 from jax-ml:pjit-lin-rule

Google-ML-Automation committed 1 year ago

030ee4a1

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

9f6dbef3

Add linearization rule for pjit_p

dougalm committed 1 year ago

b1d1dcf6

[sharding_in_types] Make argmax and argmin work with sharding_in_types. This also requires adding `reduce_p` sharding rule

yashk2810 committed 1 year ago

21f8885a

Use with_spec where possible to clean up the code a bit

yashk2810 committed 1 year ago

76356052

Adds coverage for spmd-axisname-filtering in shard_map transpose.

jkr26 committed 1 year ago

c0811c9d

Add a jaxlib at head build to the cloud-tpu-ci-nightly workflow

nitins17 committed 1 year ago

34a2f0ca

[Pallas] Deprecate dictionary compiler_params in favor of dataclass.

justinjfu committed 1 year ago

73fa0f48

[sharding_in_types] Add scan support to sharding_in_types. There are a couple of changes here

yashk2810 committed 1 year ago

355589f3

Merge pull request #25048 from jax-ml:linearization-rule-signature

Google-ML-Automation committed 1 year ago

3d79df24

Change signature of linearization rules.

dougalm committed 1 year ago

170718c8

[Pallas] Add readme page for debugging tips.

justinjfu committed 1 year ago

344d0d99

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

26443bbd

[Mosaic TPU] Fold sublane offset to indices when storing to untiled ref.

bythew3i committed 1 year ago

f899d515

Remove unneeded dependency from rocm_plugin_extension.

klucke committed 1 year ago

f3e7e682

Merge pull request #25009 from jakevdp:keyarray

Google-ML-Automation committed 1 year ago

c1ae13b0

[pallas:mosaic_gpu] `emit_pipeline` now correctly supports `BlockSpec`s in GMEM

superbobry committed 1 year ago

1efef6bf

Fix false positive `debug_nans` error caused by NaNs that are properly handled in `jax.scipy.stats.gamma`

dfm committed 1 year ago

96c01299

Merge pull request #25034 from gnecula:poly_state

Google-ML-Automation committed 1 year ago

e707edea

[pallas] Add more test cases for Triton bitcast_convert_type lowering rule.

petebu committed 1 year ago

2178ed2f

[mgpu] Pointwise op can handle LHS splats.

cperivol committed 1 year ago

1d2dc17e

Merge pull request #25018 from jakevdp:update-array-api

Google-ML-Automation committed 1 year ago

b1b1ad62

Fix cron schedule to run past minute 0 every 2nd hour

nitins17 committed 1 year ago

1e6654a0

Merge pull request #25015 from barnesjoseph:add-google-sans

Google-ML-Automation committed 1 year ago

73352677

[JAX] Ignore xla_gpu_experimental_autotune_cache_mode when calculating module hash.

Google-ML-Automation committed 1 year ago

bf0150bb

Newer Older