PR #271 CI: 03/11/25 upstream sync

fix unit tests to not use fmha rewriter

Cjkkkk committed 1 year ago

fix member access to packed CUDA struct

h-vetinari committed 1 year ago

fix Initializer protocol

Qwlouse committed 1 year ago

Fixed printing order of results in jax.debug.print documentation.

Sai-Suraj-27 committed 1 year ago

Remove duplicate JAX version 0.4.37 heading in changelog

rajasekharporeddy committed 1 year ago

Improve error messages when input argument resolution fails in custom_* APIs.

dfm committed 1 year ago

Merge pull request #26765 from Qwlouse:patch-1

Google-ML-Automation committed 1 year ago

Fix detection of epath

NeilGirdhar committed 1 year ago

Update array-api-tests to latest commit

jakevdp committed 1 year ago

Merge pull request #26912 from dfm:resolve-args-error-message

Google-ML-Automation committed 1 year ago

Merge pull request #26916 from jakevdp:update-array-api-tests

Google-ML-Automation committed 1 year ago

[JAX] Disable msan and asan for the profiler test running on nvidia gpu

Google-ML-Automation committed 1 year ago

jax.numpy ndim/shape/size: deprecate non-array input

jakevdp committed 1 year ago

Internal config change

Google-ML-Automation committed 1 year ago

[Mosaic GPU] Add lowering for `log`, and a fast path using log2.

Google-ML-Automation committed 1 year ago

Update setup.py to automatically pick up libtpu patch releases

skye committed 1 year ago

Prepare for JAX release 0.5.2 (patch release over 0.5.1)

skye committed 1 year ago

Add functionality to allow promoting RC wheels during release

Google-ML-Automation committed 1 year ago

Merge pull request #26901 from NeilGirdhar:etils

Google-ML-Automation committed 1 year ago

Merge pull request #26641 from jakevdp:jnp-ndim

Google-ML-Automation committed 1 year ago

Update all uses of `@tsl//third_party` to `@xla//third_party`

Google-ML-Automation committed 1 year ago

Update references to JAX's GitHub repo

Google-ML-Automation committed 1 year ago

Make sure concat + vmap of sharded input and replicated input works properly.

Google-ML-Automation committed 1 year ago

Merge branch 'release/0.5.2' into main

skye committed 1 year ago

Update version number after 0.5.2 release

skye committed 1 year ago

[mgpu] Non-vector untiled stores for tiling layouts.

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

Attempt 2 at landing custom_vjp.optimize_remat using custom_dce.

Google-ML-Automation committed 1 year ago

Merge pull request #26627 from Cjkkkk:remove_fmha_rewriter

Google-ML-Automation committed 1 year ago

Remove a TSAN suppression.

hawkinsp committed 1 year ago

Merge pull request #26691 from h-vetinari:packed

Google-ML-Automation committed 1 year ago

Removed unused import

superbobry committed 1 year ago

Merge pull request #26938 from superbobry:maint-2

Google-ML-Automation committed 1 year ago

[Mosaic GPU][NFC] Refactor MMA SMEM descriptor creation

Google-ML-Automation committed 1 year ago

Merge pull request #26934 from hawkinsp:tsan

Google-ML-Automation committed 1 year ago

[Mosaic GPU] Add support for small tiles for (WG)MMA LHS

Google-ML-Automation committed 1 year ago

support e2m1fn

wenscarl committed 1 year ago

Reverts 342cb7b99a09180472823a33c7cdad8a8db77875

Google-ML-Automation committed 1 year ago

Update `.bazelrc` with Apple CC toolchain changes.

Google-ML-Automation committed 1 year ago

[Mosaic GPU] Remove support for large tiles on Blackwell

Google-ML-Automation committed 1 year ago

add pmap axes hints

lockwo committed 1 year ago

Fix ambiguous cpu definition for JAX wheels.

Google-ML-Automation committed 1 year ago

vmap mismatch size error message: handle *args

garymm committed 1 year ago

Fix build rule for free-threaded python builds.

Google-ML-Automation committed 1 year ago

Merge pull request #26952 from garymm:vmap-arg

Google-ML-Automation committed 1 year ago

[Pallas] Adds a simple dynamic race detector for TPU interpret mode.

Google-ML-Automation committed 1 year ago

Set `USERPROFILE` for Windows builds to fix CI issue.

Google-ML-Automation committed 1 year ago

Add a note about uneven sharding and with_sharding_constraint. Fixes https://github.com/jax-ml/jax/issues/26946

Google-ML-Automation committed 1 year ago

Just use `jit` as the string in error messages instead of `jit` and `pjit` based on resource_env. This is to start deprecating the need for `with mesh` and replace it with `use_mesh(mesh)`.

Google-ML-Automation committed 1 year ago

[Pallas/Mosaic GPU] Enable `ops_test` for Mosaic GPU.

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

[pallas] Add support for `pl.dot` with `int8` inputs.

Google-ML-Automation committed 1 year ago

[pallas:mosaic_gpu] Added WG lowering rule for `lax.bitcast_convert_type_p`

Google-ML-Automation committed 1 year ago

Build JAX wheels instead of installing it from the source repository

Google-ML-Automation committed 1 year ago

Minor bug fixes in error checking

Google-ML-Automation committed 1 year ago

[Mosaic GPU] Fix `scf.ForOp` lowering to put lowered ops at the right place.

Google-ML-Automation committed 1 year ago

Open source TPU-friendly ragged paged attention kernel.

Google-ML-Automation committed 1 year ago

Prevent tracer leaks in scipy.special.expn

jakevdp committed 1 year ago

Updates the Colocated Python's serialization (and deserialization) implementation to utilize the recently added support for string arrays.

Google-ML-Automation committed 1 year ago

Merge pull request #26927 from skye:merge_release

Google-ML-Automation committed 1 year ago

Auto complete specs in a sharding if aval.ndim > len(sharding.spec) with `None`. So that for a 2D input, P('data') continues to work.

Google-ML-Automation committed 1 year ago

[Mosaic TPU] Support fp8 upcast to f32

Google-ML-Automation committed 1 year ago

roofline: Support computing flops for unary ops.

Google-ML-Automation committed 1 year ago

Fix lint

wenscarl committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

Dedup non-ref constants closed in cond branch functions.

Google-ML-Automation committed 1 year ago

Merge pull request #26950 from lockwo:Owen/add-pmap-typehint

Google-ML-Automation committed 1 year ago

[Mosaic GPU] Add support for tiling stores to refs using small tiling

Google-ML-Automation committed 1 year ago

Update ad.backward_pass to support non-linear functions of constants.

dfm committed 1 year ago

Add a divisibility check so that we make sure that sharding evenly divides the shape (until this restriction is lifted) to make sure we don't create bad shardings.

Google-ML-Automation committed 1 year ago

[Mosaic GPU] Add a new layout to help with transposing WGMMA results

Google-ML-Automation committed 1 year ago

[pallas:mosaic_gpu] `copy_smem_to_gmem` now allows skipping `cp.async.commit_group`

Google-ML-Automation committed 1 year ago

[Mosaic TPU] Enable all conversions involving fp8 types on TPUv5+

Google-ML-Automation committed 1 year ago

[Mosaic GPU] Add a better explanation for the transposed layout

Google-ML-Automation committed 1 year ago

[mgpu] Debug print arrays.

Google-ML-Automation committed 1 year ago

[sharding_in_types] Fix a bug where `empty_array` in scan was created with the wrong spec when `unroll > 1`.

Google-ML-Automation committed 1 year ago

Add the len(arg) to the error message for static_argnums

Google-ML-Automation committed 1 year ago

[JAX] Fix api_benchmark broken by https://github.com/jax-ml/jax/pull/26569

Google-ML-Automation committed 1 year ago

[mutable-arrays] read values should have the same explicit sharding as ref

mattjj committed 1 year ago

shard_map and other fixes to direct-linearize

mattjj committed 1 year ago

[scan] when num_trips==0, don't generate weird size-zero reshapes

mattjj committed 1 year ago

Merge pull request #26979 from mattjj:26936

Google-ML-Automation committed 1 year ago

Merge pull request #27001 from mattjj:yash-scan

Google-ML-Automation committed 1 year ago

Merge pull request #26427 from mattjj:direct-linearize-fixes

Google-ML-Automation committed 1 year ago

Support MHA in ragged paged attention for packed type

Google-ML-Automation committed 1 year ago

fixups that we meant to include in #26427

mattjj committed 1 year ago

Merge pull request #27005 from mattjj:direct-linearize-fixes-2

Google-ML-Automation committed 1 year ago

[direct-linearize] fix name_stack bugs

mattjj committed 1 year ago

Merge pull request #27010 from mattjj:direct-linearize-fixes-3

Google-ML-Automation committed 1 year ago

[Mosaic TPU][Pallas] Add pl.reciprocal

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

A few more fixes for debug_info tests with direct_linearize.

dfm committed 1 year ago

Merge pull request #27006 from dfm:more-direct-lin-debug-info

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

Update stateful-computations.md

garymm committed 1 year ago

[Pallas/Mosaic GPU] Add lowerings/layout inference for all the necessary conversion ops when using Warpgroup semantics.

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

[pallas:mosaic_gpu] Added support for math functions to the WG lowering

Google-ML-Automation committed 1 year ago

Avoid using array operations within lax.py operations.

dfm committed 1 year ago

Move (most) jaxlib linalg custom call registration into JAX.

Google-ML-Automation committed 1 year ago

jax CI: 03/11/25 upstream sync #271 Merged

jax
CI: 03/11/25 upstream sync
#271

Merged