ROCm/jax

Pull Requests Commits

Remove 6.2.4 build

charleshofer committed 264 days ago

6985f0d9

Trivial change for CI

charleshofer committed 265 days ago

9de75a25

charleshofer committed 265 days ago

ea55e59d

Add clang config file

charleshofer committed 265 days ago

98baf09f

Move clang gcc path options to config file

charleshofer committed 265 days ago

248e6386

Make Clang use manylinux C++ standard library

charleshofer committed 268 days ago

3b4a7b02

Account for versioned clang binaries

charleshofer committed 269 days ago

16d737b0

Fix typo in the error message

yashk2810 committed 270 days 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 270 days ago

48fcf028

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

yashk2810 committed 270 days ago

ef84e9d3

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

vfdev-5 committed 270 days ago

ba0879a3

[JAX] Remove calls to xla_computation_to_mlir_module.

hawkinsp committed 270 days ago

8cd2843d

kaixih committed 270 days ago

8993c0e0

kaixih committed 270 days ago

56d13e0d

Enable public doc for scaled dot

kaixih committed 270 days ago

bfe79232

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

bythew3i committed 270 days ago

185d65fb

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

yashk2810 committed 270 days ago

8a6bfd6d

Split weakref_lru_cache into its own extension.

hawkinsp committed 270 days ago

d094562b

Finalize deprecation of `ffi_call` with inline arguments.

dfm committed 270 days ago

686144b0

Remove reexports of ml_dtypes types from xla_client.py.

hawkinsp committed 270 days ago

879b72a6

Added `jax.no_tracing` to the API docs

superbobry committed 270 days ago

8d711063

Simplify handling of type stubs in nanobind extension rules.

hawkinsp committed 270 days ago

7fd3a075

Clarify jax.make_jaxpr docstring

j-towns committed 270 days ago

872a43d7

jnp.linalg: add symmetrize_input argument & docs

jakevdp committed 270 days ago

6965601d

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

jakevdp committed 270 days ago

84210a32

mattjj committed 270 days ago

689f7661

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

mattjj committed 270 days ago

67df04c6

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

bythew3i committed 270 days ago

4e094523

[shard-map] add while_map rep rule

mattjj committed 270 days ago

54e8df41

[mutable-arrays] limit implicit ref_swap dtype promotion

mattjj committed 270 days ago

4b4f8285

Older