1 Release published by 1 person

• jax-v0.7.0 JAX v0.7.0

  published Jul 22, 2025

128 Pull requests merged by 15 people

• Support _get_physical_tpu_mesh for TPU7X.

  #30490 merged Jul 24, 2025

• [pallas:triton] Migrated example kernels to plgpu.{load,store}

  #30476 merged Jul 24, 2025

• [Mosaic][SC] Add custom assembly format to tpu.enqueue_indirect_dma

  #30293 merged Jul 24, 2025

• [Mosaic GPU] Explicitly set kernel_name on both Pallas and plain Mosaic GPU codepaths.

  #30453 merged Jul 24, 2025

• [jax:custom_partitioning] Allow factors for non-batching dimensions to

  #30291 merged Jul 24, 2025

• Skip Pallas and Mosaic GPU tests that don't fit on RTX 6000 PRO

  #30258 merged Jul 24, 2025

• [JAX] Move 7x mesh creation code to OSS.

  #30478 merged Jul 24, 2025

• Make tiling argument of Layout class public and the property too. So basically it's s/_tiling/tiling/g.

  #30477 merged Jul 24, 2025

• lax.zeros_like_array is no longer public

  #30429 merged Jul 24, 2025

• Remove local_config_nvshmem repository from JAX WORKSPACE file.

  #30471 merged Jul 24, 2025

• Replace the docker image with pre-installed CUDA with a clean Docker image.

  #30473 merged Jul 24, 2025

• [Mosaic] Allow vector::Extract for non-32 bits vector result.

  #30371 merged Jul 24, 2025

• [bazel] Create jax/_src/BUILD and move _src build declarations

  #30452 merged Jul 24, 2025

• Remove --@local_config_nvshmem//:include_nvshmem_libs usage and update mosaic GPU code to conform to the new _solib folder name.

  #30448 merged Jul 24, 2025

• [pallas] Forked load and store into triton and tpu

  #30384 merged Jul 24, 2025

• [Mosaic GPU] Add missing namespace specified

  #30456 merged Jul 24, 2025

• [pallas] Moved atomic APIs to jax.experimental.pallas.triton

  #30427 merged Jul 24, 2025

• [build] avoid top-level dependencies in jax/_src

  #30461 merged Jul 24, 2025

• Refactor: Add carry support to nd_loop

  #30237 merged Jul 24, 2025

• Allow fusing all stride == 1 to support jax.lax.slice_in_dim in @fuser.fuse

  #30425 merged Jul 24, 2025

• before iterating, copy dict that can be popped from via ARC / object deletion

  #30465 merged Jul 24, 2025

• add psend/precv to auto-generated docs

  #30464 merged Jul 24, 2025

• [Mosaic GPU] Fix some typos in docs

  #30447 merged Jul 23, 2025

• Refactoring HTTIJ and surrounding pages, take 2

  #30420 merged Jul 23, 2025

• avoid top-level imports in source_info_util

  #30446 merged Jul 23, 2025

• Test B200 against CUDA 12.8 only

  #30383 merged Jul 23, 2025

• Add Hermetic C++ Toolchains for Linux x86_64 builds.

  #29672 merged Jul 23, 2025

• [Mosaic GPU] Query amount of shared memory programmatically

  #30257 merged Jul 23, 2025

• [jax:ffi] Expose build_ffi_lowering_function to give user access to the lowering function that returns an Op, not just Op results.

  #30434 merged Jul 23, 2025

• [Pallas TPU] Remove use_sreg_for_state flag

  #30421 merged Jul 23, 2025

• Remove out_sharding from scatter_add because we now have a scatter sharding rule

  #30443 merged Jul 23, 2025

• Use the correct tracing context when creating zeros in linearize_from_jvp

  #30440 merged Jul 23, 2025

• Implement explicit mode sharding rule for scatter

  #30377 merged Jul 23, 2025

• #sdy Add unit tests making sure we fallback to GSPMD if the original model was exported with GSPMD but loaded with Shardy.

  #30190 merged Jul 23, 2025

• [jax] Export ir_attribute in jax.interpreters.mlir

  #30430 merged Jul 23, 2025

• Fix typo in doc.

  #30431 merged Jul 23, 2025

• [Mosaic] NFC: Just use ExtFOp and TruncFOp for bf16.

  #30398 merged Jul 23, 2025

• [Mosaic GPU] Simplify type constraints in MGPU dialect.

  #30397 merged Jul 23, 2025

• [Mosaic GPU] Add support for warp shuffles with elements wider than 32-bit

  #30387 merged Jul 23, 2025

• Fix a bug where doing a CPU jax.Array -> TPU jax.Array transfer was not setting the layout correctly on the output TPU array.

  #30424 merged Jul 23, 2025

• [Mosaic] Add tpu.wait_indirect_dma assembly format

  #30169 merged Jul 23, 2025

• Extend numpy.size to accept multiple axes.

  #30132 merged Jul 23, 2025

• Avoid top-level imports in jax/_src/internal_test_util/*

  #30402 merged Jul 23, 2025

• Add portpicker to the list of common requirements and update requirements_lock_3_13_ft.txt and requirements_lock_3_14_ft.txt files.

  #30372 merged Jul 23, 2025

• Avoid top-level imports in jax/_src/extend

  #30403 merged Jul 23, 2025

• Remove JAX_IFRT_VERSION_NUMBER guards after the release of 0.7.0

  #30419 merged Jul 22, 2025

• #sdy Delete JAX test configs enabling Shardy.

  #30418 merged Jul 22, 2025

• Add LRU cache on get_process_index_and_count.

  #30417 merged Jul 22, 2025

• Avoid top-level imports in jax/_src/test_multiprocess

  #30415 merged Jul 22, 2025

• Remove vmap and pmap mentions in the rst section title.

  #30416 merged Jul 22, 2025

• Add smap to the API docs on the website. Fixes https://github.com/jax-ml/jax/issues/30407

  #30411 merged Jul 22, 2025

• Remove jaxlib_extension_version, ifrt_version guards now that 0.7.0 release is out

  #30406 merged Jul 22, 2025

• Reverts d86857a9c8cdcae8d27e6a371daceb0ac3f53e28

  #30408 merged Jul 22, 2025

• Release 0.7.0

  #30404 merged Jul 22, 2025

• Refactor "How to think in JAX" and surrounding pages

  #29541 merged Jul 22, 2025

• Avoid top-level imports in jax/_src/test_util.py

  #30400 merged Jul 22, 2025

• [jax/BUILD] avoid _src dependencies on top-level jax packages

  #30394 merged Jul 22, 2025

• [JAX] Make util.lru_cache(..., trace_context_in_key=False) always cache

  #30392 merged Jul 22, 2025

• Enforce the shape of MMA operands and result in their types.

  #30390 merged Jul 22, 2025

• [pallas:mosaic_gpu] Do not use lax.zeros_like_array

  #30385 merged Jul 22, 2025

• [pallas] pl.loop now accepts step=

  #30274 merged Jul 22, 2025

• Update deprecation schedule for mlir.custom_call, attributes of xla_client, and xla_bridge.get_backend.

  #30386 merged Jul 22, 2025

• [pallas:mosaic_gpu] async_load_tmem no longer accepts an idx

  #30382 merged Jul 22, 2025

• Rollback XLA archival change, current automation does not calculate the shasum correctly

  #30379 merged Jul 22, 2025

• [Mosaic GPU][NFC] Implement inline_mgpu and custom_primitive for WG semantics.

  #30147 merged Jul 22, 2025

• Add visibility to //jax:pretty_printer

  #30373 merged Jul 22, 2025

• Fix jax caching bug. The memory space selection logic should take into account the memory space of the requested sharding.

  #30374 merged Jul 22, 2025

• write docs about controlling array layouts

  #30368 merged Jul 22, 2025

• Remove SafeNumDevices now that all shardings are implemented in C++

  #30375 merged Jul 22, 2025

• Loosen test tolerance for tests/linalg_sharding_test.py::LinalgShardingTest::test_batch_axis_sharding_jvp.

  #30370 merged Jul 22, 2025

• Fix handling of empty arrays in ufunc.reduce/accumulate

  #29248 merged Jul 21, 2025

• Fix mesh_cast + new rng keys usage

  #30367 merged Jul 21, 2025

• [Pallas] Only run TPU interpret mode tests on CPU.

  #30366 merged Jul 21, 2025

• If the input aval mesh is empty, use the mesh of out_sharding passed to with_sharding_constraint on the output aval. We only use the mesh and not the pspec. This leads to more tracing cache hits as demonstrated in the test.

  #30360 merged Jul 21, 2025

• Fix py_import dependencies.

  #30365 merged Jul 21, 2025

• [Pallas][Mosaic GPU] Fix racy test by having each warpgroup work on non-overlapping data.

  #30354 merged Jul 21, 2025

• [JAX] Purge more caches when clearing backends

  #30324 merged Jul 21, 2025

• Migrate third party uses of jax.lib.xla_bridge.get_backend to jax.extend.backend.get_backend.

  #30359 merged Jul 21, 2025

• [Pallas] Make semantics of BlockSpec.memory_space in emit_pipeline consistent with pallas_call.

  #30284 merged Jul 21, 2025

• jnp.fft.fftfreq: fix regression in complex dtype support

  #30356 merged Jul 21, 2025

• [jex] Add proto <-> xla_client.HloSharding and proto -> xla_client.OpSharding` APIs to jex.

  #30357 merged Jul 21, 2025

• Fetch XLA archive using the GitHub API endpoint instead of the web link

  #30358 merged Jul 21, 2025

• Parametrize build system on CUDA major version

  #28968 merged Jul 21, 2025

• Allow setting transfer_size for cross host transfers.

  #30355 merged Jul 21, 2025

• gumbel distribution implementation

  #29343 merged Jul 21, 2025

• Automated g4 rollback of changelist 784947368.

  #30353 merged Jul 21, 2025

• Consider this example:

  #30350 merged Jul 21, 2025

• [Mosaic] Add store canonicalization for an expand reshape->store fusion

  #30333 merged Jul 21, 2025

• Add note about direct linearize to changelog

  #30351 merged Jul 21, 2025

• Update editorconfig to the Python line length of 88

  #30348 merged Jul 21, 2025

• Repair initializers not matching Initializer protocol

  #29561 merged Jul 21, 2025

• test_multiprocess: only initialize GPUs that exist

  #30248 merged Jul 21, 2025

• #sdy Add section on Shardy migration issues about all meshes being the same and custom_partitioning.

  #30347 merged Jul 21, 2025

• update use_shardy_partitioner description.

  #30345 merged Jul 21, 2025

• improve AOT input sharding/layout mismatch error message

  #30344 merged Jul 21, 2025

• Fix creation of sharding rules for fused_attention_stablehlo.

  #30343 merged Jul 20, 2025

• [JAX] Enable Shardy by default in JAX.

  #30302 merged Jul 20, 2025

• [mutable-arrays] remat discharge rule

  #29370 merged Jul 19, 2025

• [mutable-arrays] basic shard_map + mutable arrays support

  #30340 merged Jul 19, 2025

• Add JAX tests for deadlock verifier

  #30264 merged Jul 19, 2025

• Fix CI config for non-rbe CUDA tests with py_import.

  #30329 merged Jul 19, 2025

• #sdy Fix forward of making JAX changes so we can fall back to GSPMD in JAX export if the loaded module was lowered for GSPMD.

  #30336 merged Jul 19, 2025

• [mosaic] tpu_custom_call.CostEstimate is now a typed dict

  #30303 merged Jul 19, 2025

• Reverts 7c394f41574b1f669cbb6a47fd03bf1925a19444

  #30331 merged Jul 19, 2025

• [Pallas/TPU] Interpret mode: Don't use source_info context manager if it's None

  #30326 merged Jul 19, 2025

• Use operand.aval.sharding in convert_element_type's transpose rule

  #30330 merged Jul 19, 2025

• Remove the assert from scatter_addsub_transpose_rule. Fixes https://github.com/jax-ml/jax/issues/30327. We don't need to use the out_sharding in transpose rule anyways that's bound to the primitive.

  #30328 merged Jul 19, 2025

• Fixed a bug previously hidden by jax_use_direct_linearize=True

  #30325 merged Jul 19, 2025

• Move ShapeDtypeStruct to core.py to break circular deps

  #30323 merged Jul 18, 2025

• Create workflow for testing bazel cuda non rbe + py import

  #28860 merged Jul 18, 2025

• Replace jax.stages.OutInfo with jax.ShapeDtypeStruct

  #30318 merged Jul 18, 2025

• #sdy condition jax export tests on compatibility version when using shardy

  #30312 merged Jul 18, 2025

• [Mosaic GPU] Fix barrier indexing in tests.

  #30307 merged Jul 18, 2025

• Exit the CUDA test job early if not all wheels can be downloaded.

  #30316 merged Jul 18, 2025

• Add eval_shape as a method of Traced.

  #30314 merged Jul 18, 2025

• #sdy fix-forward jaxlib compatibility issue

  #30311 merged Jul 18, 2025

• If you are doing jax.eval_shape on a jitted function, call jitted_fn.eval_shape(*args) internally as a shortcut.

  #30313 merged Jul 18, 2025

• Use direct linearization by default

  #30262 merged Jul 18, 2025

• [Scaled Matmul] add a sharding rule and fix custom partition method

  #30250 merged Jul 18, 2025

• [JAX] add a missing sharding rule in ffi.md and ffi.ipynb

  #30305 merged Jul 18, 2025

• Add migration doc for Shardy in JAX, in preperation for making Shardy the default in the upcoming JAX release.

  #30301 merged Jul 18, 2025

• [Mosaic GPU] Forbid the assignment of splat layouts to non-splat constants.

  #30280 merged Jul 18, 2025

• [Mosaic GPU] Prevent vector.loads from being assigned splat layouts.

  #30278 merged Jul 18, 2025

• [Mosaic GPU] Introduce constraints in the equational layout inference.

  #30275 merged Jul 18, 2025

• #sdy add missing include in py_client.cc for Shardy.

  #30300 merged Jul 18, 2025

• [Mosaic GPU] Add lowering for tmem_alloc and tmem_dealloc.

  #30223 merged Jul 18, 2025

• [Pallas/Mosaic GPU] Add a deviceless export test for Pallas/Mosaic GPU.

  #30298 merged Jul 18, 2025

• [Mosaic] Add store canonicalization for an expand reshape->store fusion

  #30254 merged Jul 18, 2025

63 Pull requests opened by 10 people

• [JAX] Disable Shardy in `JaxExportTest` and `CompatTest` if jaxlib version is before 0.7.0.

  #30304 opened Jul 18, 2025

• [JAX] Disable Shardy in JAX export if jaxlib version is before 0.7.0.

  #30306 opened Jul 18, 2025

• Test PR for runner

  #30315 opened Jul 18, 2025

• Update ragged_dot kernels to use new GroupInfo for persistence

  #30317 opened Jul 18, 2025

• Reverts dd59b47c07caa777f57637107379374e7906de12

  #30319 opened Jul 18, 2025

• Reverts c9700e637550b6404e85aeae1ff4eb207e1f2d76

  #30320 opened Jul 18, 2025

• flip flag to check OSS tests

  #30321 opened Jul 18, 2025

• [Pallas:MGPU] Expose TCGEN05_TMEM_NATIVE_COL

  #30322 opened Jul 18, 2025

• Add aval out to pull_block_spec signature in fusible dtype

  #30332 opened Jul 19, 2025

• Automated Code Change

  #30334 opened Jul 19, 2025

• fix core.Tracer inheritance (just to see what breaks)

  #30341 opened Jul 19, 2025

• [mutable-arrays] flip JAX_MUTABLE_ARRAY_CHECKS=True by default

  #30342 opened Jul 20, 2025

• Fix numerical bugs in the gradients of sigmoid/logistic, tanh, expm1.

  #30346 opened Jul 21, 2025

• Accelerate deprecation for jax.lib.xla_bridge.get_backend.

  #30349 opened Jul 21, 2025

• Add more Bazel tests to Nightly/Release job.

  #30361 opened Jul 21, 2025

• Bump fonttools from 4.51.0 to 4.59.0

  #30362 opened Jul 21, 2025

• Bump fsspec from 2024.5.0 to 2025.7.0

  #30363 opened Jul 21, 2025

• Bump hypothesis from 6.102.4 to 6.136.1

  #30364 opened Jul 21, 2025

• [Mosaic] Allow matrix-vector dot.

  #30369 opened Jul 21, 2025

• Bump libtpu version before release.

  #30380 opened Jul 22, 2025

• #sdy Delete JAX test configs enabling Shardy.

  #30381 opened Jul 22, 2025

• Finalize deprecation of xla_bridge.get_compile_options.

  #30388 opened Jul 22, 2025

• Finalize deprecation of xla_extension.

  #30389 opened Jul 22, 2025

• Refactor Bazel CPU RBE and Bazel GPU Non-RBE and add more Bazel tests to Nightly/Release job.

  #30393 opened Jul 22, 2025

• #sdy Delete JAX test configs enabling Shardy.

  #30395 opened Jul 22, 2025

• [Mosaic] Lower exp2 as math.exp2.

  #30396 opened Jul 22, 2025

• First prototype of explicit async scheduling feature

  #30399 opened Jul 22, 2025

• Remove meaningless sdpa util test

  #30405 opened Jul 22, 2025

• [Mosaic][SC] Declare `tpu.scatter_store`. Stores values of a vector to arbitrary locations in memory.

  #30409 opened Jul 22, 2025

• Reverts 8061b583e29fb62a39c876048e9a0109ca69beec

  #30410 opened Jul 22, 2025

• Calculate roofline peak HBM bytes incrementally.

  #30412 opened Jul 22, 2025

• Roofline: stop printing entire jaxpr when erroring.

  #30413 opened Jul 22, 2025

• [Pallas][Mosaic GPU] Support pytree in/out specs in the warp specialized pipeline.

  #30414 opened Jul 22, 2025

• Improve jax.nn.standardize numerical stability

  #30428 opened Jul 23, 2025

• [pallas:mgpu] Add batching rule for `plgpu.kernel`.

  #30433 opened Jul 23, 2025

• Fix breakage caused by cl/785610645.

  #30435 opened Jul 23, 2025

• Expose GPU tracing knobs for 3P

  #30436 opened Jul 23, 2025

• [Mosaic GPU] Allow tcgen05.mma with different swizzles for A and B

  #30438 opened Jul 23, 2025

• [RPA] Add random seed and disable unaligned num_heads for RPA test

  #30444 opened Jul 23, 2025

• Implement PjRtClient::MakeCrossHostReceiveBuffers and PjRtBuffer:CopyToRemoteDevice in the PjRt C API.

  #30449 opened Jul 23, 2025

• Adding tests for the LİSHT activation function

  #30451 opened Jul 23, 2025

• [JAX][Mosaic] Add `llvm` namespace to `cast` op in mosaic_gpu.cc.

  #30454 opened Jul 23, 2025

• Add JAX tests for deadlock verifier

  #30455 opened Jul 23, 2025

• Add option to enable coordination service client to recover.

  #30457 opened Jul 23, 2025

• Add option to enable coordination service client to recover.

  #30458 opened Jul 23, 2025

• Add option to enable TFRT GPU client in C API plugin.

  #30459 opened Jul 23, 2025

• Add environment variable to enable TFRT GPU client in GPU plugin.

  #30460 opened Jul 23, 2025

• [JAX] Remove enable_empty_arrays config flag. Behavior is now always True.

  #30463 opened Jul 24, 2025

• Correctly apply global scale in nvfp4 quantization.

  #30467 opened Jul 24, 2025

• Remove default behavior of out_sharding = operand.sharding for scatter set/add as it may involve collectives.

  #30468 opened Jul 24, 2025

• Create composite action for downloading wheels from GCS

  #30472 opened Jul 24, 2025

• Add support for deserializing xplanes to Jaxlib

  #30475 opened Jul 24, 2025

• [mosaic] Added a verifier for `tpu.unpack_subelements`

  #30479 opened Jul 24, 2025

• Enable libnvptxcompiler and libnvjitlink support in OSS by default.

  #30480 opened Jul 24, 2025

• [Mosaic GPU] Predicate `tmem_alloc` and `tmem_dealloc` at the warp level.

  #30482 opened Jul 24, 2025

• [Mosaic GPU] Add `async_tmem_load` and `async_tmem_store` ops to the Mosaic MLIR dialect.

  #30483 opened Jul 24, 2025

• [Mosaic GPU] Implement lowering rules for `async_tmem_load` and `async_tmem_store`.

  #30484 opened Jul 24, 2025

• fix: typo in direct_linearize_migration.md

  #30485 opened Jul 24, 2025

• [XLA:GPU] update deadlocks verifier behavior to get computations from the scheduler

  #30486 opened Jul 24, 2025

• core.Tracer shouldn't inherit from basearray.Array

  #30487 opened Jul 24, 2025

• Pipe incarnations through C plugin.

  #30488 opened Jul 24, 2025

• Pallas/Mosaic support host output.

  #30489 opened Jul 24, 2025

• Add Bazel TPU tests in JAX presubmit and postsubmit.

  #30492 opened Jul 25, 2025

19 Issues closed by 6 people

• Build error: jaxlib 0.7.0, GCC 13.2.0

  #30437 closed Jul 24, 2025

• Segfault while Performing All-to-All Collective Operation on 8xH100 SXM5 (Shard + Swapaxes + Shard)

  #30335 closed Jul 24, 2025

• jax.grad precision: float32 gradients of bfloat16 weights

  #30337 closed Jul 23, 2025

• Fusing the optimizer step into the backward pass

  #30338 closed Jul 23, 2025

• [jax._src.stages.ArgInfo] does not declare `__len__`, breaking `jit(f).lower()`

  #30441 closed Jul 23, 2025

• jax.smap is not in the online documentation

  #30407 closed Jul 22, 2025

• [HPC] Missing libdevice.so.10

  #22590 closed Jul 22, 2025

• Debugging `jaxlib.xla_extension.XlaRuntimeError: INTERNAL: RET_CHECK failure (external/xla/xla/service/cpu/ir_emitter.cc:1144)`

  #22484 closed Jul 22, 2025

• indexing and sharding

  #13632 closed Jul 22, 2025

• CuSolver: Switch to 64 bit api to allow for eigh on matrices > than 26732x26732

  #23413 closed Jul 22, 2025

• [How?] Debugging crash in jax.jit function when running under pytest-xdist

  #10242 closed Jul 22, 2025

• [RFE] Add support for distributed CPU-backend mode

  #11182 closed Jul 22, 2025

• Failed build: CI - with Numpy&#x2F;Scipy nightly wheels (nightly)

  #30055 closed Jul 22, 2025

• Add Gumbel distribution to scipy.stats

  #29319 closed Jul 21, 2025

• jax.numpy.fft.fftfreq no longer supports complex dtype

  #30287 closed Jul 21, 2025

• //tests/mosaic:gpu_test_distributed_gpu and //tests/pallas:gpu_pallas_distributed_test_gpu fail on single-GPU machines

  #29513 closed Jul 21, 2025

• `jax.tree_util.tree_map` fails when a registered pydantic object which has been copied using `deep=True`

  #30299 closed Jul 19, 2025

• Could you add the support of the new optimizer: Muon

  #30309 closed Jul 19, 2025

• [sharding-in-types] setting global mesh+complex tensors+linear solve = problems

  #30327 closed Jul 19, 2025

14 Issues opened by 14 people

• jax.scipy.optimize not accessible in 0.7.0

  #30491 opened Jul 25, 2025

• Segmentation fault (core dumped) on H100

  #30481 opened Jul 24, 2025

• Improve jax.config.update experience in interactive environments

  #30474 opened Jul 24, 2025

• Failed to determine best cudnn convolution algorithm for large kernel conv with large stride.

  #30469 opened Jul 24, 2025

• BCOO is not compatible with `jax.lax.scan`

  #30466 opened Jul 24, 2025

• `visualize_array_sharding` unable to show device name on some shapes

  #30462 opened Jul 23, 2025

• `jax.nn.standardize` returns `nan` when variance is lower than `-epsilon`

  #30426 opened Jul 23, 2025

• Not Implemented functions in Pallas

  #30423 opened Jul 23, 2025

• Revisiting `custom_jvp` with `pure_callback`

  #30401 opened Jul 22, 2025

• Pallas fails to write to output ref when using bfloat16 on TPU v3-8

  #30391 opened Jul 22, 2025

• Profiling tool shows no GPU processes despite GPU Util being as expected.

  #30378 opened Jul 22, 2025

• jet and equinox.nn.MLP

  #30352 opened Jul 21, 2025

• Inconsistent Reduction precision in backwards computation

  #30310 opened Jul 18, 2025

• XlaRuntimeError: INTERNAL: Failed to capture gpu graph: an unsupported value or parameter was passed to the function

  #30297 opened Jul 18, 2025

25 Unresolved conversations

Sometimes conversations happen on old items that aren’t yet closed. Here is a list of all the Issues and Pull Requests with unresolved conversations.

• added solve_sylvester and accompanying tests

  #28810 commented on Jul 23, 2025 • 5 new comments

• [CI] Add bazel TPU presubmit testing

  #29660 commented on Jul 25, 2025 • 4 new comments

• [Pallas/TPU] Add option to allow skipping the device barrier

  #30289 commented on Jul 19, 2025 • 0 new comments

• Implement performance optimized w8a8 pallas kernel

  #30268 commented on Jul 18, 2025 • 0 new comments

• Remove `local_config_nvshmem` repository and corresponding macros.

  #30267 commented on Jul 23, 2025 • 0 new comments

• Support core axis index in the `device_id` dict for async copy and semaphore.

  #30243 commented on Jul 24, 2025 • 0 new comments

• Implementing LİSHT

  #30218 commented on Jul 23, 2025 • 0 new comments

• [jaxprs] Hoist large constants as arguments during lowering

  #30180 commented on Jul 24, 2025 • 0 new comments

• Construct XLA GPU client with coordination service client.

  #29872 commented on Jul 23, 2025 • 0 new comments

• Bump fsspec from 2024.5.0 to 2025.5.1

  #29011 commented on Jul 21, 2025 • 0 new comments

• Global mutable arrays not supported in `custom_jvp`

  #27470 commented on Jul 25, 2025 • 0 new comments

• Out of GPU memory

  #27354 commented on Jul 24, 2025 • 0 new comments

• support jax.numpy.linalg.eig eigenvector derivatives

  #2748 commented on Jul 24, 2025 • 0 new comments

• vmap of cond's predicate results in select, leading to unexpected compute/memory use

  #8409 commented on Jul 23, 2025 • 0 new comments

• [sharding-in-types] `jnp.reshape` not figuring out a reshape output sharding when it should

  #29235 commented on Jul 22, 2025 • 0 new comments

• [pallas] Does not support initialising complex values

  #30084 commented on Jul 22, 2025 • 0 new comments

• Add custom derivative for scipy.special.hyp2f1

  #30195 commented on Jul 22, 2025 • 0 new comments

• Error with `jax.custom_batching.sequential_vmap` in the `jax.ensure_compile_time_eval` context

  #29996 commented on Jul 22, 2025 • 0 new comments

• jax.numpy.sort performance regression

  #30296 commented on Jul 22, 2025 • 0 new comments

• Support autodiff of Eigendecomposition with repeated eigenvalues

  #669 commented on Jul 21, 2025 • 0 new comments

• custom jvp + transpose raises UnexpectedTracerError

  #29948 commented on Jul 21, 2025 • 0 new comments

• Json parse exceptions (and others) in perfetto traces

  #30290 commented on Jul 20, 2025 • 0 new comments

• feature request: sparse jacobian and sparse hessians

  #1032 commented on Jul 19, 2025 • 0 new comments

• Segmentation fault when calling exported choleksy on CPU

  #29610 commented on Jul 18, 2025 • 0 new comments

• "ComputeCallSignature failed" in verbose logs investigating performance loss.

  #30270 commented on Jul 18, 2025 • 0 new comments