ROCm/jax

Pull Requests Commits

alekstheod committed 11 hours ago

e4710085

Fix plugin init script

alekstheod committed 13 hours ago

75d1d883

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

Google-ML-Automation committed 22 hours ago

0a77c18d

Fix `triu` and `tril` with explicit mode

yashk2810 committed 1 day ago

ecaa24e6

[Mosaic:GPU] Disable test in OSS since it hangs

PatriosTheGreat committed 1 day ago

90940729

Merge pull request #35383 from mattjj:hijax-check

Google-ML-Automation committed 1 day ago

a705bda5

Merge pull request #35649 from jakevdp:test-util-annotation

Google-ML-Automation committed 1 day ago

1519ef86

Eliminate redundant sanitizer toolchains and simplify the ASAN pipeline

yuriivcs committed 1 day ago

69393a1e

[Mosaic TPU] Support transposing last two dims for booleans.

yueshengys committed 1 day ago

8db058b0

Reverts 3ff9e7e2e1072a01273a99276999143d51a0f06f

Google-ML-Automation committed 1 day ago

2e5faacb

[hijax][custom_vjp3] burndown

mattjj committed 1 day ago

7716804e

Delete core.abstractify.

hawkinsp committed 1 day ago

02cdd1de

Improve safe_zip error messages to show argument lengths

danielsuo committed 1 day ago

a8ef9a93

Reverts ca6fda28429a4f4a17caeb95e4e3f45a14669b01

WindQAQ committed 1 day ago

aa405930

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

hawkinsp committed 1 day ago

8399d0fd

[typing] improve annotations for jax/_src/test_util

jakevdp committed 1 day ago

2f1aaf80

Fix test failure after LLVM update.

hawkinsp committed 1 day ago

537f3046

Add DontWant and DidntWant (the best names ever)

yashk2810 committed 1 day ago

66dbacad

Add function name and source info to custom_vjp bwd error message

mattjj committed 1 day ago

3ff9e7e2

Merge pull request #35613 from jakevdp:array-operators

Google-ML-Automation committed 1 day 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 1 day ago

19e90ae0

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

rdyro committed 1 day ago

dc446a12

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

apaszke committed 1 day ago

c0751945

Merge pull request #35184 from olupton:multimem-error

Google-ML-Automation committed 1 day ago

aeba845c

Simplify input validation for jax.Array operator overloads

jakevdp committed 1 day ago

fe5292b6

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

WindQAQ committed 1 day ago

46df13a2

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

apaszke committed 1 day ago

2e8984a9

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

dimitar-asenov committed 1 day ago

f7d5f970

Replace multi_weakref_lru_cache with a C++ implementation.

hawkinsp committed 1 day ago

f2f6aa93

Merge pull request #35633 from hawkinsp:docbuild

Google-ML-Automation committed 1 day ago

fee7f6b2

Older