jax-ml/jax issues and pull requests

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

Pull Request - State: open - Opened by copybara-service[bot] about 12 hours ago

#24893 - Fix pgle test handling of directory removal.

Pull Request - State: open - Opened by copybara-service[bot] about 14 hours ago

#24892 - [Mosaic TPU] Add general tpu.vector_store and support masked store.

Pull Request - State: open - Opened by copybara-service[bot] about 18 hours ago

#24891 - [Mosaic TPU] Support relayout for mask vector

Pull Request - State: open - Opened by copybara-service[bot] about 19 hours ago

#24890 - Implement lax.pad in Pallas.

Pull Request - State: open - Opened by copybara-service[bot] about 19 hours ago

#24889 - Don't override --xla_tpu_use_enhanced_launch_barrier if explicitly set

Pull Request - State: closed - Opened by skye about 20 hours ago
Labels: pull ready

#24888 - Missing annotations

Issue - State: open - Opened by PerilousApricot about 21 hours ago - 4 comments
Labels: bug

#24887 - [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.

Pull Request - State: open - Opened by copybara-service[bot] about 21 hours ago

#24886 - Fix typos: Change 'arugments' to 'arguments'.

Pull Request - State: closed - Opened by carlosgmartin about 21 hours ago
Labels: pull ready

#24885 - Avoid repeatedly rebuilding a tuple in issubdtype.

Pull Request - State: closed - Opened by copybara-service[bot] about 23 hours ago

#24884 - Use device kind for positional sharding device set name.

Pull Request - State: open - Opened by copybara-service[bot] about 24 hours ago

#24883 - doc: link directly to installation on the main page

Pull Request - State: closed - Opened by jakevdp about 24 hours ago
Labels: documentation, pull ready

#24881 - Add a shard map replication rule for ffi_call

Pull Request - State: open - Opened by dfm 1 day ago
Labels: pull ready

#24880 - [Pallas] Increase test coverage of pl.dot.

Pull Request - State: open - Opened by copybara-service[bot] 1 day ago

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

Pull Request - State: closed - Opened by copybara-service[bot] 1 day ago

#24878 - Avoid using a contextmanager in Primitive.bind.

Pull Request - State: closed - Opened by copybara-service[bot] 1 day ago

#24877 - Add a shard map replication rule for cond_p

Pull Request - State: closed - Opened by dfm 1 day ago
Labels: pull ready

#24876 - [AutoPGLE] Fix test after pjrt cache refactoring

Pull Request - State: closed - Opened by copybara-service[bot] 1 day ago

#24875 - ⚠️ Nightly upstream-dev CI failed ⚠️

Issue - State: open - Opened by github-actions[bot] 1 day ago - 1 comment
Labels: CI

#24874 - Add square_p

Pull Request - State: closed - Opened by pearu 1 day ago - 3 comments
Labels: pull ready

#24873 - Stop using generators for linear_util transformations

Pull Request - State: closed - Opened by dougalm 1 day ago
Labels: pull ready

#24872 - array.at.set is incredibly slow for complex128 dtype

Issue - State: open - Opened by chrisrothUT 1 day ago - 3 comments
Labels: bug

#24871 - Add numpy.put_along_axis.

Pull Request - State: open - Opened by carlosgmartin 2 days ago - 8 comments
Labels: pull ready

#24870 - Dedent your yields!

Pull Request - State: closed - Opened by copybara-service[bot] 2 days ago

#24869 - AOT errors when setting compile options

Issue - State: open - Opened by man2machine 2 days ago
Labels: bug

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

Pull Request - State: closed - Opened by copybara-service[bot] 2 days ago

#24867 - Unsupported type in metal PJRT plugin with rng_bit_generator

Issue - State: open - Opened by dmarro89 2 days ago
Labels: bug, Apple GPU (Metal) plugin

#24866 - CUDA_ERROR_SYSTEM_NOT_READY: system not yet initialized

Issue - State: open - Opened by carlosgmartin 2 days ago - 1 comment
Labels: bug

#24865 - finalize deprecation of jax.clear_backends

Pull Request - State: open - Opened by jakevdp 2 days ago
Labels: pull ready

#24864 - jnp.logaddexp2: simplify implementation

Pull Request - State: open - Opened by jakevdp 2 days ago
Labels: pull ready

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

Pull Request - State: closed - Opened by copybara-service[bot] 2 days ago

#24862 - Move the Control Flow section of Sharp Bits into its own tutorial.

Pull Request - State: open - Opened by emilyfertig 2 days ago
Labels: documentation

#24861 - [JAX/version] Support adding custom version suffix.

Pull Request - State: open - Opened by yliu120 2 days ago

#24860 - Add GPU overflow bugfix (#24846) to changelog.

Pull Request - State: closed - Opened by dfm 2 days ago
Labels: pull ready

#24858 - cond error by using Tracer as a condition

Issue - State: closed - Opened by 2300504237 2 days ago - 3 comments
Labels: needs info

#24857 - pjit handling of static_argnames is broken when using "dynamic" python strings.

Issue - State: closed - Opened by etarassov 2 days ago - 5 comments
Labels: bug, needs info

#24856 - Remove GPU test with unreasonably large memory footprint.

Pull Request - State: closed - Opened by copybara-service[bot] 2 days ago

#24855 - [pallas:triton] Fix reshape lowering with scalar output shape.

Pull Request - State: closed - Opened by copybara-service[bot] 2 days ago

#24854 - Make sure compilation_cache.is_cache_used always returns a bool

Pull Request - State: closed - Opened by hurryabit 2 days ago
Labels: pull ready

#24853 - [JAX/distributed] Small clean ups on warning emissions in jax.distributed.

Pull Request - State: open - Opened by yliu120 2 days ago

#24852 - [Mosaic] Fix mask creation for packed sublanes

Pull Request - State: closed - Opened by copybara-service[bot] 3 days ago

#24851 - Fix pgle profiling, broken in previous change.

Pull Request - State: closed - Opened by copybara-service[bot] 3 days ago

#24850 - Explicitly raise an error if more than 65535 channels are created

Pull Request - State: closed - Opened by copybara-service[bot] 3 days ago

#24849 - 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).

Pull Request - State: closed - Opened by copybara-service[bot] 3 days ago

#24848 - 7x7 `nnx.Conv` using `float32` parameter dtype overflows(?) to `nan` when sharded

Issue - State: closed - Opened by joshhansen 3 days ago - 1 comment
Labels: bug

#24847 - Add wraparound for 2x2x2 v5p

Pull Request - State: closed - Opened by copybara-service[bot] 3 days ago

#24846 - Fix overflow error in GPU batched linear algebra kernels.

Pull Request - State: closed - Opened by copybara-service[bot] 3 days ago

#24843 - Wrong determinant results for large batch

Issue - State: closed - Opened by ChenAo-Phys 3 days ago - 5 comments
Labels: bug

#24842 - jax.scipy.linalg.toeplitz: support implicit batching

Pull Request - State: closed - Opened by jakevdp 3 days ago
Labels: pull ready

#24841 - Bump actions/cache from 4.1.1 to 4.1.2

Pull Request - State: closed - Opened by dependabot[bot] 3 days ago
Labels: pull ready, dependencies, github_actions

#24840 - Bump hypothesis from 6.102.4 to 6.118.7

Pull Request - State: open - Opened by dependabot[bot] 3 days ago
Labels: dependencies, python

#24839 - [Mosaic GPU] Only run tests requiring sm90a on Hopper

Pull Request - State: closed - Opened by andportnoy 3 days ago
Labels: pull ready

#24838 - Deepcopy of pjit functions failing (or not supported)

Issue - State: open - Opened by jlperla 3 days ago
Labels: bug

#24836 - [XLA:GPU] Skip small tile sizes for sparse gemms on Ampere as well. Enable the JAX test again that has been failing.

Pull Request - State: closed - Opened by copybara-service[bot] 3 days ago

#24835 - [pallas:triton] Simplify reshape lowering rule.

Pull Request - State: closed - Opened by copybara-service[bot] 3 days ago

#24834 - Add Pallas Triton lowering for jax.lax.bitcast_convert_type.

Pull Request - State: closed - Opened by copybara-service[bot] 3 days ago

#24833 - [pallas:triton] Simplify lowering code. `BlockInfo` is now always present for memory refs.

Pull Request - State: closed - Opened by copybara-service[bot] 3 days ago

#24832 - [XLA:GPU] Change `assert` to `CHECK` in Triton sparsity extensions.

Pull Request - State: closed - Opened by copybara-service[bot] 3 days ago

#24831 - Improvements in the handling of tracing debug info

Pull Request - State: open - Opened by gnecula 3 days ago
Labels: pull ready

#24830 - support q offset w.r.t k/v in flash_attention function

Issue - State: open - Opened by yamingx 3 days ago - 2 comments
Labels: enhancement

#24829 - Bug in the latest libtpu nightly release

Issue - State: open - Opened by knyazer 4 days ago - 1 comment
Labels: bug

#24828 - Fix buggy and confusing logic in the C++/pjit caching path.

Pull Request - State: closed - Opened by dougalm 4 days ago
Labels: pull ready

#24826 - potential jax 0.4.35 release issue?

Issue - State: open - Opened by haohuanw 4 days ago - 1 comment
Labels: bug

#24824 - ⚠️ Nightly upstream-dev CI failed ⚠️

Issue - State: closed - Opened by github-actions[bot] 4 days ago - 1 comment
Labels: CI

#24823 - [shape_poly] Fix the handling of jvp(lax.sort)

Pull Request - State: closed - Opened by gnecula 4 days ago
Labels: pull ready

#24822 - Cleanup more remnants of the jax.experimental.host_callback

Pull Request - State: closed - Opened by gnecula 4 days ago
Labels: pull ready

#24821 - TPU with sharding: grpc initialization failure

Issue - State: closed - Opened by knyazer 5 days ago - 2 comments
Labels: bug

#24820 - Pattern match dot algorithm spec to preset name

Pull Request - State: open - Opened by dfm 5 days ago
Labels: pull ready

#24818 - [shape_poly] Remove caching for the symbolic shape evaluator

Pull Request - State: closed - Opened by gnecula 5 days ago
Labels: pull ready

#24817 - [Mosaic GPU] Ensure that lowering `InitializeBarrierOp` preserves the result's type.

Pull Request - State: closed - Opened by copybara-service[bot] 6 days ago

#24816 - [Mosaic GPU] Add `base_pointer` argument to `InitializeBarrierOp`.

Pull Request - State: closed - Opened by copybara-service[bot] 6 days ago

#24814 - jnp.bincount: support boolean inputs

Pull Request - State: closed - Opened by jakevdp 6 days ago
Labels: pull ready

#24813 - bincount rejects bool

Issue - State: closed - Opened by carlosgmartin 6 days ago
Labels: bug

#24812 - Skip test_jnp_einsum_grad_y_pallas on gpu due to ooms

Pull Request - State: closed - Opened by copybara-service[bot] 6 days ago

#24810 - Specialize ufunc.reduce for monoidal binary ufuncs.

Pull Request - State: closed - Opened by jakevdp 6 days ago
Labels: pull ready

#24809 - [Pallas] Add a cost estimator for Pallas/JAX functions.

Pull Request - State: closed - Opened by copybara-service[bot] 6 days ago

#24806 - Internal: create decorators for defining ufuncs

Pull Request - State: closed - Opened by jakevdp 6 days ago
Labels: pull ready

#24805 - [Mosaic GPU] Add CUPTI profiler alongside events-based implementation

Pull Request - State: open - Opened by andportnoy 6 days ago

#24800 - Allow more output storage types for some dot algorithms.

Pull Request - State: closed - Opened by copybara-service[bot] 6 days ago

#24799 - pre-commit run --all failed for ruff and mypy

Issue - State: closed - Opened by apivovarov 6 days ago - 1 comment
Labels: bug

#24798 - Flaky test tests/fft_test.py::FftTest::testFftfreq5

Issue - State: open - Opened by apivovarov 6 days ago - 2 comments
Labels: bug

#24797 - [AutoPGLE] Explicitly ignore host callback pointers

Pull Request - State: open - Opened by copybara-service[bot] 6 days ago

#24796 - Check mem_stats key in test_lax_full_like_efficient

Pull Request - State: open - Opened by apivovarov 6 days ago

#24795 - Sparse reshape throws error when `n_dense>0` and some target dimension has size 1

Issue - State: open - Opened by cherrywoods 6 days ago - 1 comment
Labels: bug

#24794 - `jax.lax.DotAlgorithm` generates an unnecessary conversion on H100

Issue - State: closed - Opened by andportnoy 6 days ago - 3 comments
Labels: bug

#24793 - Fix parallel pgle-tests execution.

Pull Request - State: closed - Opened by copybara-service[bot] 6 days ago

#24792 - Update some outdated syntax in FFI tutorial

Pull Request - State: closed - Opened by dfm 6 days ago
Labels: pull ready

#24791 - [Mosaic GPU] Make sure to free the cloned MLIR module when debugging

Pull Request - State: closed - Opened by copybara-service[bot] 6 days ago

#24790 - Add new CI scripts for running Pytests

Pull Request - State: open - Opened by copybara-service[bot] 6 days ago

#24789 - Disable lax_test on ARM in Google's internal CI.

Pull Request - State: closed - Opened by copybara-service[bot] 6 days ago

#24788 - Add commands to run Bazel GPU (non-RBE) jobs

Pull Request - State: open - Opened by copybara-service[bot] 6 days ago

#24787 - Nightly builds on Mac ARM fail complex function numerical tests

Issue - State: open - Opened by hawkinsp 6 days ago - 2 comments
Labels: bug

#24786 - Don't perform size 0 slices into scipy rotations.

Pull Request - State: closed - Opened by hawkinsp 6 days ago - 1 comment
Labels: pull ready

#24785 - Add new CI scripts for building JAX artifacts

Pull Request - State: open - Opened by copybara-service[bot] 6 days ago

#24784 - Disable the paged_attention test on TPU v5p.

Pull Request - State: closed - Opened by copybara-service[bot] 6 days ago

#24783 - Disable for_loop_test on TPU v5p.

Pull Request - State: closed - Opened by copybara-service[bot] 6 days ago

#24782 - [pallas] Allow user to pass 64-bit indices to `pl.{load,store,...}`.

Pull Request - State: open - Opened by copybara-service[bot] 6 days ago

#24781 - [Mosaic TPU] Support 1D concat: set implicit_dim to kSecondMinor to treat 1D (N,) as (1, N) and then tile it as (1, 128)

Pull Request - State: open - Opened by copybara-service[bot] 6 days ago

#24780 - [mgpu] Broadcast the LHS fragmented array if it is splat.

Pull Request - State: open - Opened by copybara-service[bot] 6 days ago

#24779 - Added suppressions for pytype --none-is-not-bool.

Pull Request - State: closed - Opened by copybara-service[bot] 6 days ago

GitHub / jax-ml/jax issues and pull requests