jax

CI: 11/14/24 upstream sync

#142

Merged

Internal: create decorators for defining ufuncs

Allow setting default_device with platform names.

jax.scipy.linalg.toeplitz: support implicit batching

Update some outdated syntax in FFI tutorial.

Add an example FFI call to demonstrate the use of global state.

Add an option to specify mock GPU topology

Add GPU overflow bugfix (#24846) to changelog.

Merge pull request #24860 from dfm:overflow-changelog

Fix pgle profiling, broken in previous change.

Merge pull request #24806 from jakevdp:ufunc-decorator

Specialize ufunc.reduce for monoidal binary ufuncs.

Change to internal dead code elimination. Now the functions in `dce_rules` are responsible for checking if the equation has no used outputs or effects, and behaving appropriately in that case (which usually means eliminating said equation).

[Mosaic] Fix mask creation for packed sublanes

[mosaic_gpu] The profiler now uses FFI calls for creating events and computing elapsed time

Merge pull request #24792 from dfm:ffi-tutorial-outdated

Update XLA dependency to use revision

Change the assumed width of the bool packing in the early-lowering checks in pallas

[Mosaic] Add parameter names to tpu.sem_signal and add tests

Dedent your yields!

Merge pull request #24842 from jakevdp:batched-toeplitz

Add new CI scripts for running Bazel CPU presubmits

[XLA] Make our LLVM usage more googley

Fix parallel pgle-tests execution.

Add a shard map replication rule for cond_p.

Avoid using a contextmanager in Primitive.bind.

Fix bug where the Python wrapper to ParseArguments() didn't intern the static argnames strings, causing false mismatches when searching for static arguments.

[AutoPGLE] Fix test after pjrt cache refactoring

Merge pull request #24877 from dfm:cond-rep-rule

Add square_p

Add Pallas Triton lowering for jax.lax.bitcast_convert_type.

When caching is enabled, also enable XLA caching features as well

Merge pull request #24069 from sergachev:cudnn_fusion_test_a100

Merge pull request #22899 from trevor-m:cache

doc: link directly to installation on the main page

Fix typos: Change 'arugments' to 'arguments'.

Merge pull request #24874 from pearu:pearu/square_p

Update XLA dependency to use revision

Stop using generators for linear_util transformations.

Don't override --xla_tpu_use_enhanced_launch_barrier if explicitly set

Merge pull request #24883 from jakevdp:doc-installation

Merge pull request #24886 from carlosgmartin:fix_typos

Avoid repeatedly rebuilding a tuple in issubdtype.

Merge pull request #24810 from jakevdp:bitwise-reduce

Merge pull request #24873 from jax-ml:no-gen-linear-util

Merge pull request #24889 from skye:barrier

Merge pull request #23805 from dfm:ffi-examples-state

[pallas:mosaic_gpu] Ported two pipelining optimizations to `emit_pipeline`

Make the jaxpr for jnp.pad in "constant" mode more succinct.

Don't use an out-of-line lowering for integer_pow for small powers.

Merge pull request #23374 from jaro-sevcik:mock-topology-config

Fixed typo in asan.yaml

finalize deprecation of jax.clear_backends

Merge pull request #24899 from vfdev-5:patch-2

Fix pgle test handling of directory removal.

Merge pull request #24865 from jakevdp:dep-clear-backends

[Mosaic][TPU] Omit short circuiting of relayout (we should always relayout!) and implement product mismatch case for where we relayout from replicated to offset, and the number of vregs changes.

[sharding_in_types] Support shard_map with sharding in types. Right now only full manual mode is supported.

Merge pull request #24751 from Stella-S-Yan:feature/default_device_str

Merge pull request #24896 from hawkinsp:pad

Merge pull request #24897 from hawkinsp:ipow

Merge remote-tracking branch 'origin/rocm-main' into ci-upstream-sync-25_1

Add commit to see if it triggers CI