ROCm/jax

Pull Requests Commits

Update documentation and add setup.py pypi bindings

Ruturaj4 committed 1 year ago

b7343a70

Remove dead code after minimum jaxlib version bump to v0.4.36.

hawkinsp committed 1 year ago

79318a08

Merge pull request #25320 from ROCm:gh-9948-fix-kernel-build-upstream

Google-ML-Automation committed 1 year ago

cc258f5f

Ensured that JAX type checks under pytype on Python 3.12

superbobry committed 1 year ago

1ac6b762

Merge pull request #25338 from carlosgmartin:fix_numpy_linalg_matrix_norm_ord_type_annotation

Google-ML-Automation committed 1 year ago

5a1c4c57

Activate Tridiagonal Reduction to XLA's FFI

Paweł Paruzel committed 1 year ago

d474feda

[Mosaic TPU] Allow downgrading the IR during serialization for forward compat

apaszke committed 1 year ago

adb2bf62

[pallas] Add `DotAlgorithmPreset` note to CHANGELOG.

chr1sj0nes committed 1 year ago

a94474d0

[pallas:triton] Add support for `DotAlgorithmPreset` `precision` arguments to `dot`.

chr1sj0nes committed 1 year ago

3ec55c77

Fix type annotation for numpy.linalg.matrix_norm argument 'ord'.

carlosgmartin committed 1 year ago

efa35ea9

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

70623255

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

ad00ee1e

Temporarily allow bfloat16 dot algorithms on CPU.

dfm committed 1 year ago

1f4d184a

Support transfer guard in broadcast_one_to_all(). Fixes https://github.com/jax-ml/jax/issues/25325

danijar committed 1 year ago

861115ad

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

baedb62b

Add a flag to enable detailed timestamped logging of subprocess commands.

nitins17 committed 1 year ago

83c64b23

Merge pull request #25082 from nouiz:doc_cuda_malloc_async

Google-ML-Automation committed 1 year ago

b6499e21

Document cudaMallocAsync as an experimental feature.

nouiz committed 1 year ago

a13b618c

[Pallas] Add support for run_state to cost estimator.

justinjfu committed 1 year ago

641a1d53

Merge pull request #25205 from jburnim:jburnim_swap_fix

Google-ML-Automation committed 1 year ago

72df8e0c

Don't look for CUDA files when building the ROCm wheel

charleshofer committed 1 year ago

0c6b967e

[Pallas MGPU] Disable XLA:GPU autotuning in attention tests

apaszke committed 1 year ago

eda7506d

[Pallas MGPU] Use multiple k/v_consumed_barriers in the attention kernel

apaszke committed 1 year ago

8b656206

Fix error when swapping a ref with a trivial indexing transform.

jburnim committed 1 year ago

af501356

[mosaic_gpu] Emit a slightly more informative error message in `FragmentedArray._pointwise`

superbobry committed 1 year ago

08d31d0f

[pallas:mosaic_gpu] `FragmentedArray.reduce_sum` now returns a `FragmentedArray`

superbobry committed 1 year ago

bae66000

Set -Werror=mismatched-tags on clang.

hawkinsp committed 1 year ago

fac1b1a7

Activate Schur Decomposition to XLA's FFI

Paweł Paruzel committed 1 year ago

9081e85d

Bump JAX version after release.

hawkinsp committed 1 year ago

ba626fa6

Merge pull request #25252 from gnecula:poly_power

Google-ML-Automation committed 1 year ago

9fc077a5

Older