ROCm/jax

Pull Requests Commits

RNN fix MIOpen wights layout

shurale-nkn committed 22 days ago

7fb13d9f

Merge pull request #40784 from ROCm:user/mminutoli/disable-rocm-latest-on-release

Google-ML-Automation committed 22 days ago

ecaee84b

[typing] improve type annotation for jnp.float32, jnp.int32, and other scalar types

Jake VanderPlas committed 22 days ago

568f805f

Merge pull request #40799 from hawkinsp:fixbuild

Google-ML-Automation committed 22 days ago

3b95c706

[hijax] jnp.nonzero: compute via hijax primitive

Jake VanderPlas committed 22 days ago

0e15d558

Fix build breakage.

hawkinsp committed 22 days ago

e8ded89f

[ROCm] Skip the TheRock pre-release CI legs on release runs

mminutoli committed 22 days ago

b7ab9325

Merge pull request #39953 from hawkinsp:npy21

Google-ML-Automation committed 22 days ago

6a4f54eb

[Roll forward -- Take 3] Optimize large JAX Mesh construction

yashk2810 committed 22 days ago

c69f2ec7

Drop support for NumPy 2.1.

hawkinsp committed 22 days ago

4f1e40d3

Make `jnp.ldexp` and `jnp.frexp` bit-exact for normal floating-point numbers.

hawkinsp committed 22 days ago

7fc04171

Remove version checks that are now always false:

hawkinsp committed 22 days ago

53b742ae

Measure real ULP errors in `numerics` tests (`elementary_test` and `special_test`).

hawkinsp committed 22 days ago

8fd3d2a4

Upgrade rules_cc to 0.2.25 in JAX

akuegel committed 22 days ago

41aacdc2

[hijax] jnp.searchsorted: compute via hijax primitive

Jake VanderPlas committed 22 days ago

e6c4b91b

PR #40790: Integrate 0.11.2 release branch and update version numbers.

hawkinsp committed 22 days ago

5191eb6a

[pallas:mgpu] Added support for custom derivative primitives in the lowering

superbobry committed 22 days ago

324e831c

[mosaic_gpu] `parse_indices` no longer allows `slice()` with an `ir.Value`

superbobry committed 22 days ago

eb1f8ada

api_test: Use a more future-proof way to detect the C API usage.

pizzud committed 22 days ago

0fe26af7

[Roll forward -- Take 2] Optimize large JAX Mesh construction

Google-ML-Automation committed 22 days ago

d925891a

[Pallas:Triton] Mask out-of-bounds memory accesses for unaligned blocks

akuegel committed 22 days ago

45d520ca

[Roll forward -- Take 2] Optimize large JAX Mesh construction

Google-ML-Automation committed 22 days ago

008c954b

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

Google-ML-Automation committed 22 days ago

db378d24

[emit_pipeline] Fix async prefetch for the no outputs case.

rdyro committed 23 days ago

533182d0

Merge pull request #40338 from teddytennant:partition-order-invert

Google-ML-Automation committed 23 days ago

438c014b

Fix partition, argpartition and top_k on signed integer minimums

teddytennant committed 23 days ago

5a51ddd7

[Pallas] Skip test_async_sc_tc_prefetch_vmem with libtpu < 0.0.49.

hawkinsp committed 23 days ago

5df3e770

[Pallas][SC] Remove needs_layout_passes=False from debug_print tests

Google-ML-Automation committed 23 days ago

fc7fe71f

Lower shard_map in partial explicit manual mode correctly by using the aval.sharding.spec instead of `in_specs` since the former contain more info (explicit mesh axes too) while in_specs only refer to axes passed in axis_names.

yashk2810 committed 23 days ago

ae0ca5b8

Acquire a global lock in `jax.experimental.colocated_python.serialization._deserialize` to synchronize `cloudpickle.loads` calls across multiple threads.

Google-ML-Automation committed 23 days ago

20c6952b

Older