ROCm/jax

Pull Requests Commits

Delete core.abstractify.

hawkinsp committed 32 days ago

02cdd1de

Improve safe_zip error messages to show argument lengths

danielsuo committed 32 days ago

a8ef9a93

Reverts ca6fda28429a4f4a17caeb95e4e3f45a14669b01

WindQAQ committed 32 days ago

aa405930

Replace uses of core.get_aval (deprecated) with core.typeof.

hawkinsp committed 32 days ago

8399d0fd

[typing] improve annotations for jax/_src/test_util

jakevdp committed 32 days ago

2f1aaf80

AratiGanesh committed 32 days ago

6efa9a3c

Fix test failure after LLVM update.

hawkinsp committed 32 days ago

537f3046

Add DontWant and DidntWant (the best names ever)

yashk2810 committed 32 days ago

66dbacad

Add function name and source info to custom_vjp bwd error message

mattjj committed 32 days ago

3ff9e7e2

Merge pull request #35613 from jakevdp:array-operators

Google-ML-Automation committed 32 days ago

ff10f745

Fix device_put_transpose in presence of memory spaces so that memory_space of ct matches primal. Also fix a bug in dp_transpose rule where if arg was an undefined primal but ct was a Zero, we were returning None instead of Zero.

yashk2810 committed 32 days ago

19e90ae0

Move the TPU lowering basic_ragged_dot case (reliant on a fusion) to the

rdyro committed 33 days ago

dc446a12

[Mosaic TPU] Add support for float -> unsigned int casts (arith.fptoui)

apaszke committed 33 days ago

c0751945

Merge pull request #35184 from olupton:multimem-error

Google-ML-Automation committed 33 days ago

aeba845c

Simplify input validation for jax.Array operator overloads

jakevdp committed 33 days ago

fe5292b6

[Mosaic] NFC: Refactor tpu::PackMaskOp to take > 2 sources.

WindQAQ committed 33 days ago

46df13a2

[Mosaic TPU] Add first-class support for unsigned integer -> float (arith.uitofp)

apaszke committed 33 days ago

2e8984a9

[Mosaic GPU] Add support for SMEM aliasing using WG semantics in the Pallas lowering.

dimitar-asenov committed 33 days ago

f7d5f970

Install ml_dtypes build dependencies to fix upstream numpy nightly build.

hawkinsp committed 33 days ago

42bd8d44

Replace multi_weakref_lru_cache with a C++ implementation.

hawkinsp committed 33 days ago

f2f6aa93

Merge pull request #35633 from hawkinsp:docbuild

Google-ML-Automation committed 33 days ago

fee7f6b2

Update Python version to 3.12 in documentation build.

hawkinsp committed 33 days ago

7ff8de77

[Pallas:Triton] Fix a type checker error to fix CI

apaszke committed 33 days ago

f08b2083

[Mosaic GPU][NFC] Add missing type annotations to `FragmentedArray`.

allanrenucci committed 33 days ago

a0de217d

Merge pull request #34598 from ROCm:phambinh/ci_rocm_approx_tanh

Google-ML-Automation committed 33 days ago

d9f8bf02

Merge pull request #35603 from superbobry:pyrefly

Google-ML-Automation committed 33 days ago

704d0a03

Removed blanket `type: ignore`s which are not necessary under Pyrefly

superbobry committed 33 days ago

6a8a5c8d

Merge pull request #34602 from ROCm:phambinh/ci_rocm_pallas_ops_tests

Google-ML-Automation committed 33 days ago

08d80110

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

Google-ML-Automation committed 33 days ago

0ffa1103

Properly mirror CompilationProviderOptions for jax

ermilovmaxim committed 33 days ago

e527f48c

Newer Older