ROCm/jax

Pull Requests Commits

Remove Primitive._true_bind.

hawkinsp committed 14 days ago

2b0b00d6

Add a missing export in jax.debug.

Federico Lebrón committed 14 days ago

2deba493

Merge pull request #35679 from hawkinsp:docs2

Google-ML-Automation committed 14 days ago

bc975f8e

[Mosaic GPU] Remove SwizzleTransform from MGPU

apaszke committed 14 days ago

3bfa917d

Remove dtypes.is_string_dtype, just offer dtypes.string_dtype.

hawkinsp committed 14 days ago

5c86fc03

Monkey-patch jupyter_client to avoid concurrent kernel starts.

hawkinsp committed 14 days ago

401c7a5c

[pallas:sc] Allowed `VectorSubcoreMesh` to use a subset of subcores

superbobry committed 14 days ago

653d0f77

[Mosaic GPU] Add support for atomic reductions other than add

apaszke committed 14 days ago

8f7bf88f

Change core.valid_jaxtype to allow tracers.

hawkinsp committed 14 days ago

498ef4e7

Merge pull request #35634 from hawkinsp:nnightly

Google-ML-Automation committed 14 days ago

df2ec0c8

Fix Pallas test failure after LLVM update.

hawkinsp committed 14 days ago

06ca68e3

Merge pull request #34768 from ROCm:fix_rnn_struct_encoding_upstream

Google-ML-Automation committed 14 days ago

80ed0e31

Merge pull request #35557 from gnecula:poly_clean1

Google-ML-Automation committed 14 days ago

a1ff74b3

[export] Remove code to serialize HloSharding

gnecula committed 14 days ago

5b55e64b

[Mosaic GPU] Remove workarounds for an XLA:GPU bug that's been fixed for a while now

apaszke committed 14 days ago

e19ec903

[Mosaic TPU] Allow loading of custom dialects when reserializing the IR

apaszke committed 14 days ago

c357fc9a

[Mosaic GPU] Add support for vectorized atomics

apaszke committed 14 days ago

88a1f94f

[Pallas:TPU] Use the newly added uint <-> float conversions when available.

apaszke committed 14 days ago

c7a3187b

[Pallas:TPU] Fix pow_p lowering to use uitofp for unsigned integer exponents

apaszke committed 14 days ago

bdf67038

[Mosaic:TPU][NFC] Remove unused propagateTiledLayoutsToConsumers after cl/878457195

tlongeri committed 14 days ago

2a8cc7af

Add proxy fences to `dynamic_scheduling_loop`.

chr1sj0nes committed 14 days ago

4b7e7903

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

Google-ML-Automation committed 14 days ago

0a77c18d

Fix `triu` and `tril` with explicit mode

yashk2810 committed 14 days ago

ecaa24e6

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

PatriosTheGreat committed 14 days ago

90940729

Merge pull request #35383 from mattjj:hijax-check

Google-ML-Automation committed 14 days ago

a705bda5

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

Google-ML-Automation committed 14 days ago

1519ef86

Eliminate redundant sanitizer toolchains and simplify the ASAN pipeline

yuriivcs committed 14 days ago

69393a1e

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

yueshengys committed 14 days ago

8db058b0

Reverts 3ff9e7e2e1072a01273a99276999143d51a0f06f

Google-ML-Automation committed 14 days ago

2e5faacb

[hijax][custom_vjp3] burndown

mattjj committed 14 days ago

7716804e

Newer Older