ROCm/jax

Pull Requests Commits

[ROCm] Implement RNN support

Ruturaj4 committed 1 year ago

fe68eb8b

Add a register_custom_type_id function to the GPU plugins.

dfm committed 1 year ago

a7f384cc

Merge pull request #25748 from shoyer:divmod

Google-ML-Automation committed 1 year ago

853af560

Merge pull request #25633 from dfm:move-ffi

Google-ML-Automation committed 1 year ago

56f0f953

Merge pull request #25731 from gnecula:poly_random_even

Google-ML-Automation committed 1 year ago

712bece2

Fix type signature for __divmod__

shoyer committed 1 year ago

7fb68cac

[shape_poly] Improve threefry with symbolic shapes

gnecula committed 1 year ago

bc3306c8

Merge pull request #25728 from zhenying-liu:scipy.misc

Google-ML-Automation committed 1 year ago

7997f080

Make inspect_array_sharding work without mesh context manager too.

yashk2810 committed 1 year ago

23eaf216

Merge pull request #25740 from jakevdp:remove-array-api

Google-ML-Automation committed 1 year ago

b304b9ef

Remove deprecated jax.experimental.array_api

jakevdp committed 1 year ago

c7b0d681

[Mosaic TPU] Add some elementwise canonicalizations

sharadmv committed 1 year ago

4caa263a

Remove the need for check_rep for with_sharding_constraint.

pschuh committed 1 year ago

b49ba655

fix the doc error: module 'scipy.misc' has no attribute 'face'

zhenying-liu committed 1 year ago

77c6947a

Rename pybind_extension to nanobind_extension.

hawkinsp committed 1 year ago

90d8f378

Suppress MSAN warnings from SVD that are showing up in CI.

hawkinsp committed 1 year ago

61dd0412

Merge pull request #25214 from jakevdp:einsum-optimize

Google-ML-Automation committed 1 year ago

52cc5c7f

[Mosaic TPU] Validate inserted layout in relayout-insertion pass.

bythew3i committed 1 year ago

9f842909

Merge pull request #25699 from yliu120:fix_iota

Google-ML-Automation committed 1 year ago

634b45bf

jnp.einsum: default to optimize='auto'

jakevdp committed 1 year ago

2f7204ff

bazel: export serialization.fbs for downstream usage

maxwillzq committed 1 year ago

c39e38fe

Merge pull request #25675 from jakevdp:dep-lpmn

Google-ML-Automation committed 1 year ago

74be8bd9

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

18b193cb

Deprecate scipy.special.lpmn & lpmn_values

jakevdp committed 1 year ago

245a13a3

Fixes tril/triu comments (they were flipped)

marksandler2 committed 1 year ago

6c87bf38

fix the degenerated case

yliu120 committed 1 year ago

3ff000ee

[shape_poly] Remove old non_negative support.

gnecula committed 1 year ago

e87a2a59

Disable `avxvnniint8` when building with Clang version < 19, or GCC < 13.

belitskiy committed 1 year ago

f2e210b3

Temporarily allow deprecation warnings for `scipy.special.lpmn` and `scipy.special.sph_harm`.

dfm committed 1 year ago

512d5450

Merge pull request #25715 from ROCm:ci_build_code_fixes-upstream

Google-ML-Automation committed 1 year ago

4de7794f

Older