ROCm/jax

Pull Requests Commits

[ROCm] Move ROCm Bazel presubmit workflow to 32-bit mode.

tsrw2048 committed 66 days ago

ede818c9

Pad the partitions with None if it's length is less than the `axis` in get_sharding_for_vmap

yashk2810 committed 66 days ago

20a74963

[pallas:mosaic] `pl.debug_print` now supports 32-bit floats

superbobry committed 66 days ago

64a2fcfb

Reverts 776018ae0b4f2972f15cd23673fbca817a267f7d

zvikinoza committed 66 days ago

bff1d544

Use pspec.partitions in _make_lengths_same

Google-ML-Automation committed 66 days ago

e7c30496

Expose `jax.ffi.abstract_token`

Jake VanderPlas committed 66 days ago

a6c992da

[Pallas:MGPU] Add Lane semantics lowering for `lax.concatenate`.

allanrenucci committed 66 days ago

67cae6e2

[mosaic] Enabled `use-nameloc-as-prefix` for IR dumps

superbobry committed 66 days ago

055a433b

[Mosaic GPU] Add support for concatenating `FragmentedArray`s.

allanrenucci committed 66 days ago

425c9c03

Add PDL kernel launch integration to pallas mosaic gpu

rajaanand-g committed 66 days ago

5c65d574

Fix bug where jnp.int64(1 << 40) would truncate even if explicit x64 types are allowed.

hawkinsp committed 66 days ago

f7337b26

Merge pull request #38694 from aybchan:aybchan/skip-tile-scatter4

Google-ML-Automation committed 66 days ago

5f8829ca

Merge pull request #38724 from vfdev-5:pallas-fix-_gmem_slice_and_predicate

Google-ML-Automation committed 66 days ago

3673b557

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

Google-ML-Automation committed 67 days ago

4c290b69

Reverts 52e03d80559778007f3d9a379a06ce8f003332d8

brianwa84 committed 67 days ago

ca202035

Enable HBM input buffer in `emit_pipeline` (via a simple window slice).

rdyro committed 67 days ago

2f682d05

Merge pull request #38779 from mattjj:ct-check-refaccum

Google-ML-Automation committed 67 days ago

f6237c2d

Skip a ragged_collective_test that is crashing on the released jaxlib.

hawkinsp committed 67 days ago

a723a69b

adjust internal deprecations API.

Jake VanderPlas committed 67 days ago

5e39c370

Fix use of unregistered mlir dialect in autodidax example.

hawkinsp committed 67 days ago

3d24c4a2

add correct ct check to refaccum

mattjj committed 67 days ago

9fc04c36

Merge pull request #38766 from jax-ml:dependabot/github_actions/actions/checkout-7.0.0

Google-ML-Automation committed 67 days ago

9f842b72

[Mosaic][NFC] Add helper function to get the remainder of a Value, if determinable

tlongeri committed 67 days ago

85425ccf

Merge pull request #38587 from Intel-tensorflow:ashahba/rename-oneapi-whl

Google-ML-Automation committed 67 days ago

9b6874ae

re-land #38730: Lower `jnp.empty` & `jnp.empty_like` to `lax.empty`

Jake VanderPlas committed 67 days ago

11cefe97

Merge pull request #38772 from jakevdp:fix-empty-batching

Google-ML-Automation committed 67 days ago

59c17f04

Remove fancy batching rule for lax.empty_p

jakevdp committed 67 days ago

967391f0

[doc] jnp.pad: clarify meaning of `mode="empty"`

Jake VanderPlas committed 67 days ago

ddb1b1d5

If you trace and lower with abstract mesh and abstract inputs and then do `out = lowered.compile(device_assignment=...)(real_arr)`, then make sure the `out.sharding` is concrete i.e. has real devices and is not abstract.

yashk2810 committed 67 days ago

d5bc1dd9

Lower `jnp.empty` & `jnp.empty_like` to `lax.empty`

Google-ML-Automation committed 67 days ago

184e7bf4

Older