1 Release published by 1 person

• jax-v0.7.0 JAX v0.7.0

  published Jul 22, 2025

121 Pull requests merged by 15 people

• Count jit_cpp_cache_misses by registering an event in def _pjit_call_impl_python

  #30529 merged Jul 26, 2025

• Don't define a standard_primitive definition for broadcast_in_dim_p because it also has a separate definition for abstract_eval which makes it confusing. Just define it as a normal primitive instead.

  #30525 merged Jul 26, 2025

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

  #30243 merged Jul 26, 2025

• Don't use aval's mesh if any of the axes in it as manual

  #30523 merged Jul 26, 2025

• [Mosaic GPU] Reverse the direction in which we iterate through the pallas grid to be row-major.

  #29105 merged Jul 25, 2025

• Fix AArch64 CUDA build by passing crosstool to bazel.

  #30513 merged Jul 25, 2025

• [CI] Update CI cache bazelrc settings

  #30508 merged Jul 25, 2025

• [Pallas/Fuser] Add support for trivial reduce sum

  #30514 merged Jul 25, 2025

• Upgrade the dependencies in the requirement lock files

  #30511 merged Jul 25, 2025

• [Mosaic] Lower exp2 as math.exp2.

  #30396 merged Jul 25, 2025

• Fix iteration over a global mutable object.

  #30510 merged Jul 25, 2025

• Fix _add_transpose rule to make sure that cotangent sharding matches the primal shardings.

  #30509 merged Jul 25, 2025

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

  #30444 merged Jul 25, 2025

• [Pallas TPU] Delete hbm to host DMA test when using host ref as input

  #30512 merged Jul 25, 2025

• Added a list of Pallas SC users to the BUILD file

  #30507 merged Jul 25, 2025

• Add more Bazel tests to Nightly/Release job.

  #30361 merged Jul 25, 2025

• [Mosaic GPU] Fix name of scale operands for tcgen05_mma.

  #30504 merged Jul 25, 2025

• Add fuser.py to jax wheel.

  #30502 merged Jul 25, 2025

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

  #30433 merged Jul 25, 2025

• [pallas] Do not use pl.{load,store} which will soon be removed from Pallas

  #30495 merged Jul 25, 2025

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

  #30463 merged Jul 25, 2025

• [pallas:triton] Do not import jax.experimental.pallas.triton on Windows

  #30496 merged Jul 25, 2025

• [pallas] pl.MemoryRef is no longer exported

  #30497 merged Jul 25, 2025

• [mosaic] Added a verifier for tpu.unpack_subelements

  #30479 merged Jul 25, 2025

• 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

58 Pull requests opened by 11 people

• 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

• Bump fonttools from 4.51.0 to 4.59.0

  #30362 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

• 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

• 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

• 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

• 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

• 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

• Internal changes.

  #30494 opened Jul 25, 2025

• [Mosaic GPU] Support wargroup semantic for automatic TMEM allocation.

  #30499 opened Jul 25, 2025

• Add `fuser.py` in jax wheel.

  #30500 opened Jul 25, 2025

• [Mosaic GPU] Make the vector length configurable in the TMEM_NATIVE_LAYOUT

  #30501 opened Jul 25, 2025

• Fix OSS breakage due to TPU test.

  #30503 opened Jul 25, 2025

• [pallas] Deprecated `pl.load` and `pl.store`

  #30505 opened Jul 25, 2025

• Add support for scipy.special.erfcx (scaled complementary error funct…

  #30506 opened Jul 25, 2025

• Only bind a reduce_sum_p if it has nontrivial axes

  #30515 opened Jul 25, 2025

• fix bnth layout shardy and infer_fwd_ouput_sharding

  #30516 opened Jul 25, 2025

• Reorganize the tutorials

  #30519 opened Jul 25, 2025

• Fix jaxlib ASAN github test

  #30521 opened Jul 25, 2025

• [pallas:mgpu] Switch batch dim position in `kernel` vmap rule to match new row-major ordering.

  #30522 opened Jul 25, 2025

• Fix ASAN build.

  #30524 opened Jul 26, 2025

• [pallas] Add support for fusing into a `fusible` within a `vmap`.

  #30527 opened Jul 26, 2025

17 Issues closed by 6 people

• jax.scipy.optimize not accessible in 0.7.0

  #30491 closed Jul 26, 2025

• 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

15 Issues opened by 15 people

• Additional function transformations to improve JAX usability

  #30528 opened Jul 26, 2025

• Bad link in Stateful Computations doc

  #30520 opened Jul 25, 2025

• JAX fails to detect GPU on HPC with CUDA 12.6.1 and cuDNN 8.0.4

  #30518 opened Jul 25, 2025

• SIGSEGV/SIGABRT with JAX 0.7.0

  #30517 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

26 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 25, 2025 • 0 new comments

• Remove `local_config_nvshmem` repository and corresponding macros.

  #30267 commented on Jul 23, 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 26, 2025 • 0 new comments

• Fix cache leaks for pe._cached_abstract_eval; add util.multi_weakref_lru_cache

  #30009 commented on Jul 27, 2025 • 0 new comments

• Construct XLA GPU client with coordination service client.

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

• Add metadata for CUDA and libtpu versions

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

• Bump fsspec from 2024.5.0 to 2025.5.1

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

• jax.scipy.linalg.eigh_tridiagonal() doesn't implement calculation of eigenvectors

  #14019 commented on Jul 26, 2025 • 0 new comments

• Batched version of jax.vmap

  #29587 commented on Jul 26, 2025 • 0 new comments

• support jax.numpy.linalg.eig eigenvector derivatives

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

• lax.associative_scan causes kernel crash in JAX 0.6.2 on GPU (PJRT)

  #29866 commented on Jul 25, 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

• Inconsistent Reduction precision in backwards computation

  #30310 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

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

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