ROCm/jax

Pull Requests Commits

Fix typo in the error message

yashk2810 committed 1 year ago

42751359

Rename `psum2` to `psum_invariant` and put it in `lax_parallel`. We shouldn't expose this to public API and have users use `psum` instead which will dispatch to `psum_invariant` when `check_rep=True`.

yashk2810 committed 1 year ago

f95f6a8b

Rename `pbroadcast` to `pvary` and expose it as `jax.lax.pvary`.

yashk2810 committed 1 year ago

84016bc3

Merge pull request #27804 from vfdev-5:ft-adapt-state-test-2

Google-ML-Automation committed 1 year ago

373ac2ef

[JAX] Remove calls to xla_computation_to_mlir_module.

hawkinsp committed 1 year ago

a516988b

Merge pull request #27503 from kaixih:enable_doc_scaled_dot

Google-ML-Automation committed 1 year ago

b8d9e7f4

[ragged-paged-attn] Unify kv strided load to one.

bythew3i committed 1 year ago

7b455524

Disable second order vjp tests in RunStateHypothesisTest.test_vjp if JAX_SKIP_SLOW_TESTS=true

vfdev-5 committed 1 year ago

5a340a97

Make changes to shard_map to prepare for setting `varying_axes_in_types` to True.

yashk2810 committed 1 year ago

8301c304

Split weakref_lru_cache into its own extension.

hawkinsp committed 1 year ago

b4629c23

Finalize deprecation of `ffi_call` with inline arguments.

dfm committed 1 year ago

2d44f985

Remove reexports of ml_dtypes types from xla_client.py.

hawkinsp committed 1 year ago

09fed2f6

Added `jax.no_tracing` to the API docs

superbobry committed 1 year ago

62df2e8d

Simplify handling of type stubs in nanobind extension rules.

hawkinsp committed 1 year ago

a43136b3

Merge pull request #27821 from j-towns:clarify-make-jaxpr-docstr

Google-ML-Automation committed 1 year ago

f1bcf3bb

Merge pull request #27807 from jakevdp:eigvalsh-symmetrize

Google-ML-Automation committed 1 year ago

76825a2d

Merge pull request #27836 from jakevdp:fix-repeat

Google-ML-Automation committed 1 year ago

b073e8db

Merge pull request #27809 from mattjj:26621

Google-ML-Automation committed 1 year ago

ef680634

[Mosaic TPU] Add MemRead and MemStore effects to load and store ops.

bythew3i committed 1 year ago

29cb6cd1

Merge pull request #27803 from mattjj:27644

Google-ML-Automation committed 1 year ago

03c1bf96

jnp.repeat: don't cast repeats to array, as they must be static.

jakevdp committed 1 year ago

b7d430f9

[mutable-arrays] limit implicit ref_swap dtype promotion

mattjj committed 1 year ago

4d2808c1

mattjj committed 1 year ago

ae957977

Replace references to jax.readthedocs.io with docs.jax.dev.

hawkinsp committed 1 year ago

e02faabf

[Mosaic TPU] Add support for non-32bit types in vector.extract

apaszke committed 1 year ago

b8353d1b

[pallas:mosaic_gpu] Added test for custom pretty-printing rules

superbobry committed 1 year ago

f5d73b89

[Mosaic GPU] Simplify load/store methods now that we have fewer layouts

apaszke committed 1 year ago

b926fac6

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

d6524dc4

Merge pull request #27827 from apaszke:mgpu-docs

Google-ML-Automation committed 1 year ago

aa6e7016

Add a skeleton for Pallas:Mosaic GPU documentation

apaszke committed 1 year ago

511f7820

Older