ROCm/jax

Pull Requests Commits

[pallas] Support reshape with trailing `1`s in fuser.

chr1sj0nes committed 310 days ago

aaa73657

Update XLA dependency to use revision

Google-ML-Automation committed 310 days ago

e23f6814

[pallas] Generalize support for merging dimensions in fuser reshape.

chr1sj0nes committed 310 days ago

a67fa87a

Merge pull request #29808 from gnecula:fix_close_jaxpr_cache

Google-ML-Automation committed 310 days ago

1b885e45

[caches] Fix the pe.close_jaxpr cache.

gnecula committed 310 days ago

c2707aaf

Merge pull request #30039 from gnecula:more_cache_cleanup

Google-ML-Automation committed 310 days ago

2c5bfac6

[JAX] Use experimental DCN transfer library as a fallback for PjRt-IFRT cross-host device transfers when the PjRt plugin doesn't implement the cross-host transfer APIs.

emilyfertig committed 310 days ago

e562cddc

Merge pull request #30026 from mattjj:scan-mutable-array-ad

Google-ML-Automation committed 310 days ago

97429c51

[caches] Register weakref_lru_caches with util._caches

gnecula committed 310 days ago

c216c9f1

Reverts 8d304fe7635fc450817b6abae5ef703a397e9b68

Google-ML-Automation committed 310 days ago

366095e8

[mutable-arrays] fix scan ad bugs

mattjj committed 310 days ago

b5364b70

Finalize several deprecations in jax.core and jax.lib.xla_client

Jake VanderPlas committed 310 days ago

8928510f

Merge pull request #30022 from gnecula:caches_cleanup

Google-ML-Automation committed 310 days ago

6465459f

[debug_info] Minor fix following https://github.com/jax-ml/jax/pull/30012

gnecula committed 310 days ago

7ac80f26

Fix a bug where vmapping over an explicit sharded dim on 1 mesh axis lead to a KeyError in shard_map.

yashk2810 committed 310 days ago

e8ea8107

Merge pull request #30025 from jakevdp:xla-extension-dep

Google-ML-Automation committed 310 days ago

86c6e52a

Merge pull request #30013 from jakevdp:dep-extend-ffi

Google-ML-Automation committed 310 days ago

2442d9ff

Simplify the jaxpr emitted for jnp.repeat in the case that the input dim size is 1.

hawkinsp committed 310 days ago

8d304fe7

Remove deprecated APIs jax.lib.xla_extension.ArrayImpl and XlaRuntimeError.

jakevdp committed 310 days ago

88cf87c1

[caches] Simplifications in the management of caches

gnecula committed 310 days ago

088ddb26

Finalize deprecation of jax.extend.ffi

jakevdp committed 310 days ago

cdf79a31

Finalize deprecations of jax.interpreters.xla.abstractify and pytype_aval_mappings

Jake VanderPlas committed 310 days ago

3c6903d6

[JAX] Use qualified names in tracebacks.

hawkinsp committed 311 days ago

116b7831

Merge pull request #29977 from vfdev-5:docs-typo-fix

Google-ML-Automation committed 311 days ago

c1872e23

Merge pull request #30017 from jakevdp:finalize-xla-client-deps

Google-ML-Automation committed 311 days ago

23192072

Remove a number of deprecated symbols in jax.lib.xla_client

jakevdp committed 311 days ago

f5da2482

[Mosaic GPU] Fix a broken B200 test

apaszke committed 311 days ago

b1fabd6b

[Mosaic GPU] Add support for block-scaled f4e2m1fn MMAs

apaszke committed 311 days ago

4452dc5b

Merge pull request #30012 from gnecula:fix_todo

Google-ML-Automation committed 311 days ago

35b62cd3

Merge pull request #30011 from jakevdp:fix-parallel-breakage

Google-ML-Automation committed 311 days ago

131ae7a4

Older