ROCm/jax

Pull Requests Commits

Fix ROCm builds not finding numa library

charleshofer committed 364 days ago

132f88e8

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

bchetioui committed 364 days ago

7fd32ecc

Update XLA dependency to use revision

Google-ML-Automation committed 364 days ago

b6da46ec

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

shraiysh committed 364 days ago

cb2eb157

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

yashk2810 committed 364 days ago

76dec382

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

Google-ML-Automation committed 1 year ago

02505fa7

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

ayaka14732 committed 1 year ago

988a1208

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

Google-ML-Automation committed 1 year ago

aceae84f

[Pallas] Increase tolerance in PallasOutOfBoundsInterpretTest.

jburnim committed 1 year ago

802cb33b

Merge pull request #26840 from rajasekharporeddy:testbranch1

Google-ML-Automation committed 1 year ago

d5587972

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

Google-ML-Automation committed 1 year ago

b8590816

[Pallas/Fuser] Add select_n push rule

sharadmv committed 1 year ago

81dde225

Merge pull request #27038 from jakevdp:vmap-sentinel

Google-ML-Automation committed 1 year ago

261e6e5f

Merge pull request #26977 from jakevdp:fix-expn

Google-ML-Automation committed 1 year ago

c942b0fe

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

sharadmv committed 1 year ago

87272fbe

Rename `dot_with_no_batch_dims_saveable` to `dots_with_no_batch_dims_saveable` for internal consistency

Google-ML-Automation committed 1 year ago

affe2e73

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

pravnar committed 1 year ago

b6d4fe53

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

Google-ML-Automation committed 1 year ago

64beebbf

Merge pull request #26525 from wenscarl:e2m1fn

Google-ML-Automation committed 1 year ago

18f2f19c

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

jburnim committed 1 year ago

73d20cd6

Internal: make it easier to detect the vmap sentinel

jakevdp committed 1 year ago

8ecadfdf

Remove version limit for `setuptools` dependency.

Google-ML-Automation committed 1 year ago

007fc7a6

Modify version test to consider "rc" versions as well

nitins17 committed 1 year ago

d41e9683

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

mwhittaker committed 1 year ago

5cb29949

Merge pull request #27032 from dfm:lax-dtype

Google-ML-Automation committed 1 year ago

14b215fe

Add file and zip to tsan.yaml

vfdev-5 committed 1 year ago

1bab037c

Merge pull request #26811 from dfm:direct-lin

Google-ML-Automation committed 1 year ago

ab0ce8a4

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

dimitar-asenov committed 1 year ago

d2bf034c

Merge pull request #27026 from garymm:patch-3

Google-ML-Automation committed 1 year ago

5a7ef404

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

dfm committed 1 year ago

21884d4a

Older