ROCm/jax

Pull Requests Commits

Remove 6.2.4 build

charleshofer committed 1 year ago

6985f0d9

Trivial change for CI

charleshofer committed 1 year ago

9de75a25

charleshofer committed 1 year ago

ea55e59d

Add clang config file

charleshofer committed 1 year ago

98baf09f

Move clang gcc path options to config file

charleshofer committed 1 year ago

248e6386

Make Clang use manylinux C++ standard library

charleshofer committed 1 year ago

3b4a7b02

Account for versioned clang binaries

charleshofer committed 1 year ago

16d737b0

Fix typo in the error message

yashk2810 committed 1 year ago

9af8e04f

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

48fcf028

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

yashk2810 committed 1 year ago

ef84e9d3

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

vfdev-5 committed 1 year ago

ba0879a3

[JAX] Remove calls to xla_computation_to_mlir_module.

hawkinsp committed 1 year ago

8cd2843d

kaixih committed 1 year ago

8993c0e0

kaixih committed 1 year ago

56d13e0d

Enable public doc for scaled dot

kaixih committed 1 year ago

bfe79232

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

bythew3i committed 1 year ago

185d65fb

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

yashk2810 committed 1 year ago

8a6bfd6d

Split weakref_lru_cache into its own extension.

hawkinsp committed 1 year ago

d094562b

Finalize deprecation of `ffi_call` with inline arguments.

dfm committed 1 year ago

686144b0

Remove reexports of ml_dtypes types from xla_client.py.

hawkinsp committed 1 year ago

879b72a6

Added `jax.no_tracing` to the API docs

superbobry committed 1 year ago

8d711063

Simplify handling of type stubs in nanobind extension rules.

hawkinsp committed 1 year ago

7fd3a075

Clarify jax.make_jaxpr docstring

j-towns committed 1 year ago

872a43d7

jnp.linalg: add symmetrize_input argument & docs

jakevdp committed 1 year ago

6965601d

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

jakevdp committed 1 year ago

84210a32

mattjj committed 1 year ago

689f7661

[shard-map] canonicalize rep=None to be rep={all possible axes}

mattjj committed 1 year ago

67df04c6

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

bythew3i committed 1 year ago

4e094523

[shard-map] add while_map rep rule

mattjj committed 1 year ago

54e8df41

[mutable-arrays] limit implicit ref_swap dtype promotion

mattjj committed 1 year ago

4b4f8285

Older