ROCm/jax

Pull Requests Commits

[ROCm] Remove broken legacy env vars

mrodden committed 1 year ago

9ff891df

Pallas pipeline API tweaks for more advanced pipelining patterns.

epiqueras committed 1 year ago

aa73aa00

Lower the shard count for sparse_bcoo_bcsr_test on TPU as well.

hawkinsp committed 1 year ago

85a466d7

Set PYTHONWARNINGS=error in bazel tests.

hawkinsp committed 1 year ago

70f91db8

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

d58a09fa

Support pmin and pmax in check_rep.

pschuh committed 1 year ago

5e3f7618

Add `--use_cuda_nvcc` flag to enable or disable compilation of CUDA code using NVCC.

Google-ML-Automation committed 1 year ago

6e116491

[Mosaic TPU] Support all cases for extui.

bythew3i committed 1 year ago

407dc774

[pallas:triton] Elide `program_id` calls where launch grid dimension is 1.

chr1sj0nes committed 1 year ago

be7fe878

Expose some APIs for querying trace state. This will let us move users away from

dougalm committed 1 year ago

d2ac88c1

[Pallas] Update export compatibility tests

apaszke committed 1 year ago

9114b084

Merge pull request #23806 from gspschmid:gschmid/ffi-ext-bundle

Google-ML-Automation committed 1 year ago

2c85465e

Merge pull request #23873 from hawkinsp:cumsum

Google-ML-Automation committed 1 year ago

d9a10e13

Merge pull request #23860 from jakevdp:better-error

Google-ML-Automation committed 1 year ago

26ba3063

[pallas] Allow `TransformedRef` to be passed to `pl.load` / `pl.store`, when `idx = ()`.

chr1sj0nes committed 1 year ago

8d86a047

Fix an incorrect output for jnp.cumsum.

hawkinsp committed 1 year ago

562e9e8d

Merge pull request #23744 from jakevdp:unary-ufunc

Google-ML-Automation committed 1 year ago

6860617e

[Mosaic GPU] Add support for input_output_aliases

apaszke committed 1 year ago

ae86ef16

Make jnp.negative a ufunc & add unary ufunc tests

jakevdp committed 1 year ago

6229511f

Add more informative error when static argument is passed to non-static JIT parameter

jakevdp committed 1 year ago

a44e129a

Added support for % and `select` to `mgpu.FragmentedArray`

superbobry committed 1 year ago

8196c8bf

Merge pull request #23861 from hawkinsp:warnings

Google-ML-Automation committed 1 year ago

6d351136

Remove warning filters from pyproject.toml, add local warning

hawkinsp committed 1 year ago

a0e44483

Merge pull request #23862 from jakevdp:array-api-ci

Google-ML-Automation committed 1 year ago

80cb821a

trigger array API tests for all PRs.

jakevdp committed 1 year ago

45af8742

Merge pull request #23741 from awshaichen:neuron

Google-ML-Automation committed 1 year ago

2ac1d0b8

Use `jax.make_array_from_process_local_data` API in distributed data loading doc

yashk2810 committed 1 year ago

a99ea733

Fix another errata in block-sparse kernel tutorial.

ds-hwang committed 1 year ago

e4091a67

Merge pull request #23807 from kaixih:fix_cudnn_sdpa_bwd_batcher

Google-ML-Automation committed 1 year ago

1b3d8dc4

Merge pull request #23739 from jakevdp:add-at

Google-ML-Automation committed 1 year ago

46867dc4

Older