ROCm/jax

Pull Requests Commits

Fix ROCm builds not finding numa library

charleshofer committed 341 days ago

132f88e8

[Pallas/Mosaic GPU] Explicitly disable `ops_test` on Mosaic GPU pre-Hopper.

bchetioui committed 341 days ago

7fd32ecc

Update XLA dependency to use revision

Google-ML-Automation committed 341 days ago

b6da46ec

PR #22800: Change the default value of print_operand_shape_ to false and print_large_constants_ to true.

shraiysh committed 341 days ago

cb2eb157

Under pjit the `with mesh:` context will use `use_mesh(mesh): jit` instead of tracking separately using `resource_env`.

yashk2810 committed 341 days ago

76dec382

[Pallas TPU] Remove `next_slot` SMEM tensor from pipeline emitter

Google-ML-Automation committed 341 days ago

02505fa7

Better error message when `raise_if_error()` is called within a traced context

ayaka14732 committed 341 days ago

988a1208

[Pallas] Enable skipping of floating-point operations when interpreting Pallas TPU kernels on CPU.

Google-ML-Automation committed 341 days ago

aceae84f

[Pallas] Increase tolerance in PallasOutOfBoundsInterpretTest.

jburnim committed 341 days ago

802cb33b

Merge pull request #26840 from rajasekharporeddy:testbranch1

Google-ML-Automation committed 341 days ago

d5587972

Merge pull request #26839 from Sai-Suraj-27:fix_jax.debug.print

Google-ML-Automation committed 341 days ago

b8590816

[Pallas/Fuser] Add select_n push rule

sharadmv committed 341 days ago

81dde225

Merge pull request #27038 from jakevdp:vmap-sentinel

Google-ML-Automation committed 341 days ago

261e6e5f

Merge pull request #26977 from jakevdp:fix-expn

Google-ML-Automation committed 341 days ago

c942b0fe

[Pallas/Fuser] Add debug option to fuser.fuse that prints out jaxpr

sharadmv committed 341 days ago

87272fbe

Rename `dot_with_no_batch_dims_saveable` to `dots_with_no_batch_dims_saveable` for internal consistency

Google-ML-Automation committed 342 days ago

affe2e73

Define lax.ragged_dot_general and express lax.ragged_dot in terms of it.

pravnar committed 342 days ago

b6d4fe53

Merge pull request #27035 from vfdev-5:add-file-zip-to-tsan-ci-jobs

Google-ML-Automation committed 342 days ago

64beebbf

Merge pull request #26525 from wenscarl:e2m1fn

Google-ML-Automation committed 342 days ago

18f2f19c

[Pallas] Small fix to TPU interpret mode (input_output_aliases + scalar args).

jburnim committed 342 days ago

73d20cd6

Internal: make it easier to detect the vmap sentinel

jakevdp committed 342 days ago

8ecadfdf

Remove version limit for `setuptools` dependency.

Google-ML-Automation committed 342 days ago

007fc7a6

Modify version test to consider "rc" versions as well

nitins17 committed 342 days ago

d41e9683

Warn the user if transparent huge pages aren't enabled.

mwhittaker committed 342 days ago

5cb29949

Merge pull request #27032 from dfm:lax-dtype

Google-ML-Automation committed 342 days ago

14b215fe

Add file and zip to tsan.yaml

vfdev-5 committed 342 days ago

1bab037c

Merge pull request #26811 from dfm:direct-lin

Google-ML-Automation committed 342 days ago

ab0ce8a4

[Mosaic GPU] Test the wgmma_op lowering when a is in registers.

dimitar-asenov committed 342 days ago

d2bf034c

Merge pull request #27026 from garymm:patch-3

Google-ML-Automation committed 342 days ago

5a7ef404

Move (most) jaxlib linalg custom call registration into JAX.

dfm committed 342 days ago

21884d4a

Older