jax-ml/jax issues and pull requests

#26632 - Automated Code Change

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

#26631 - canonicalize closed over values if atleast 1 mesh axis is `Manual` and all other mesh axes are `Manual` or `Auto`. This would make the canonicalization work properly with shmap partial-auto.

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

#26630 - If cur_mesh is empty and AxisTypes of Mesh passed to shmap are Explicit, then treat the axes mentioned in `auto` as explicit too. In other words, "auto" really means "don't convert to manual", ie leave the listed mesh axes as they are, whether explicit or auto type

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

#26629 - Fix jax wheel pattern in pytest workflows

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

#26628 - [sharding_in_types] Initial support for partial-auto/explicit shard_map + sharding-in-types. If the axes in `shmap(..., auto=...)` is an explicit axes in the outer mesh context, then that axis is treated as Explicit instead of Auto.

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

#26627 - Remove fmha rewriter dependency in cuDNN SDPA api

Pull Request - State: open - Opened by Cjkkkk 1 day ago

#26626 - [JAX] Implement an initial object API for colocated Python

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

#26625 - Only add new manual axes to residuals when adding axes with partial_auto.

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

#26624 - Bump ml_dtypes>=0.5.0

Pull Request - State: open - Opened by wenscarl 1 day ago - 2 comments

#26623 - Plumb layout through the creation of PjRtArrays.

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

#26622 - Fix rank promotion error in JVP of batched eigh.

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

#26621 - `jax.scipy.linalg.expm` broken under shard_map

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

#26620 - Relax one more check in partial_eval_jaxpr_nounits

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

#26619 - Add and test support for partitioning of batch dimensions in lax.linalg.

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

#26618 - Scan leads to much higher memory usage than an unrolled loop

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

#26617 - Deprecate alpha argument to trsm LAPACK kernel.

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

#26616 - pmap-ing slower with new CPU runtime

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

#26615 - fix scan doc on the `unroll` argument.

Pull Request - State: closed - Opened by froystig 2 days ago
Labels: kokoro:force-run, pull ready

#26614 - [Mosaic] Consider divisibility when doing large tiling

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

#26613 - Relax the check in `_mapped_axis_spec` to allow `()` and `None` to be treated the same

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

#26612 - Update lax.linalg.svd primitive to use registration helper functions.

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

#26611 - Update several remaining lax.linalg primitives to use registration helper functions.

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

#26610 - Mark `in_shardings` and `out_shardings` as Any for typing reasons since they can take pytrees. Fixes https://github.com/jax-ml/jax/issues/26609

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

#26609 - Pyright does not correcly infer type for `out_sharding` on `jax.jit`

Issue - State: closed - Opened by ZacCranko 2 days ago - 3 comments
Labels: bug

#26608 - `repr` or `str` method for `jax._src.config.Config` objects

Issue - State: open - Opened by ZacCranko 2 days ago
Labels: enhancement

#26607 - Add dtype dispatching to FFI example.

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

#26606 - Error out if going from `Manual` -> `Auto/Explicit` AxisTypes in the `auto_axes` and `explicit_axes` API that do `mesh_cast` implicitly.

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

#26605 - FFI: jaxlib.xla_extension.XlaRuntimeError: INVALID_ARGUMENT: Failed to serialize StableHLO;

Issue - State: closed - Opened by tilmantroester 2 days ago - 4 comments
Labels: bug

#26604 - [pallas:triton] Fix atomic min/max lowering for uint and float types

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

#26603 - [better_errors] Simplify api_util.debug_info.

Pull Request - State: open - Opened by gnecula 2 days ago - 6 comments
Labels: pull ready

#26601 - Use "common" instead of "build" for some flags in .bazelrc

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

#26597 - [Pallas][Mosaic] Support float8_e4m3b11fnuz

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

#26596 - Exclude auto axes whenever extending the axis_env via core.extend_axis_env_nd.

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

#26595 - Now that sharding_in_types config flag is True, remove the config and all the conditionals

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

#26594 - Simplify heartbeat options for PjRt distributed runtime.

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

#26593 - Expose `get_ty` aka get_aval from jax namespace

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

#26592 - Let users pass in pspecs to with_sharding_constraint when `use_mesh` is set. This is in-line with other APIs which allow pspecs like einsum, reshape, etc

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

#26591 - jax.lax: improve docs for comparison operators

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

#26590 - [Pallas] Support dynamic grids in the new TPU interpret mode

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

#26589 - roofline: Support computing flops for binary ops.

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

#26588 - [Will not merge] bug report on autodiff on partial auto nested shard_map

Pull Request - State: open - Opened by yliu120 3 days ago - 1 comment

#26587 - Update internal unop primitive helper to pass kwargs to dtype rule.

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

#26586 - Avoiding collectives in sharding vs pmap

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

#26585 - [StableHLO] Only emit StableHLO from xla_computation_to_mlir_module

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

#26584 - ClosedJaxpr: Expose inner invars, outvars, constvars.

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

#26583 - Update lax.linalg.qr primitives to use registration helper functions.

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

#26582 - Update lax.linalg.lu primitive to use registration helper functions.

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

#26581 - [pallas] Skip `OpsTest.test_concat_constant` if TPU is not available

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

#26580 - [Mosaic TPU] Support mask concat

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

#26579 - Controlling rematerialization when the model is split across multiple jitted functions

Issue - State: open - Opened by gspschmid 3 days ago - 1 comment
Labels: enhancement

#26577 - Edit type signature of lexsort in stub file to match implementation

Pull Request - State: closed - Opened by BaconBreaker 3 days ago - 1 comment

#26575 - bandpass/butterworth filters

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

#26574 - Change type signature of lexsort in stub file to match type signature in implementation

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

#26572 - Typing for Pytree

Issue - State: closed - Opened by krzysztofrusek 3 days ago - 3 comments
Labels: enhancement

#26571 - Fix typos in ffi examples

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

#26568 - Vmap slower than jax.lax.map

Issue - State: closed - Opened by dapatil211 3 days ago - 2 comments
Labels: question

#26566 - Saturating arithmetic

Issue - State: open - Opened by ivanstepanovftw 4 days ago - 3 comments

#26564 - Add jax.experimental._mini_mpmd

Pull Request - State: open - Opened by gspschmid 4 days ago - 6 comments

#26563 - Error unconditionally for jit, pjit and with_sharding_constraint if `use_mesh` and `with mesh` are used together.

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

#26561 - [pallas:mosaic_gpu] Use `emit_pipeline` for pipelining in the lowering

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

#26560 - `check_tracer_leaks` sometimes raises `IndexError` when parent is a cell from a closure.

Issue - State: open - Opened by tillahoffmann 5 days ago
Labels: bug

#26559 - jax.experimental.sparse.BCOO sum over axis does not reduce number of data entries

Issue - State: closed - Opened by GarfieldGa 5 days ago - 2 comments
Labels: question

#26558 - Typo in jax ffi examples

Issue - State: closed - Opened by 0xSooki 5 days ago - 1 comment

#26557 - Documentation missing for `drop_fields` argument of `tree_util.register_dataclass`

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

#26556 - feat(gh-13291): Add Pareto distribution functions: `cdf`, `logcdf`, `logsf`, `sf`, `ppf`, `logppf`, and corresponding tests

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

#26556 - feat(gh-13291): Add Pareto distribution functions: `cdf`, `logcdf`, `logsf`, `sf`, `ppf`, `logppf`, and corresponding tests

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

#26555 - `matrix_norm` raises `ValueError` for empty matrices

Issue - State: open - Opened by carlosgmartin 6 days ago - 3 comments
Labels: bug

#26554 - Let `seed=0` by default in `random.key`

Issue - State: open - Opened by carlosgmartin 6 days ago - 4 comments
Labels: enhancement

#26553 - Make device_put resharding on single device array input work under use_mesh. Fixes https://github.com/jax-ml/jax/issues/26552

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

#26553 - Make device_put resharding on single device array input work under use_mesh. Fixes https://github.com/jax-ml/jax/issues/26552

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

#26552 - [sharding_in_types] Unable to `device_put` an unsharded array inside `use_mesh`

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

#26552 - [sharding_in_types] Unable to `device_put` an unsharded array inside `use_mesh`

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

#26551 - [JAX] Allow pallas to accept scalar shape semaphores.

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

#26550 - jax.lax: improve docs for bitwise operators.

Pull Request - State: open - Opened by jakevdp 7 days ago
Labels: documentation

#26550 - jax.lax: improve docs for bitwise operators.

Pull Request - State: closed - Opened by jakevdp 7 days ago
Labels: documentation, pull ready

#26550 - jax.lax: improve docs for bitwise operators.

Pull Request - State: closed - Opened by jakevdp 7 days ago
Labels: documentation, pull ready

#26549 - Support error checking in auto mode outside jit

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

#26549 - Support error checking in auto mode outside jit

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

#26549 - Support error checking in auto mode outside jit

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

#26548 - Fix the type annotations and don't += a generator (it's confusing)

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

#26548 - Fix the type annotations and don't += a generator (it's confusing)

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

#26546 - [TPU][Mosaic][Easy] Add verification for AssumeMultipleOp.

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

#26545 - [CI] testing change to internal sync

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

#26545 - [CI] testing change to internal sync

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

#26544 - Add version check to jaxlib plugin imports.

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

#26544 - Add version check to jaxlib plugin imports.

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

#26543 - Start adding primitive registration helper functions to lax.linalg.

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

#26543 - Start adding primitive registration helper functions to lax.linalg.

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

#26543 - Start adding primitive registration helper functions to lax.linalg.

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

#26542 - Fix breakage in indexing refactor

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

#26542 - Fix breakage in indexing refactor

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

#26541 - Fix the type annotations and don't += a generator (it's confusing)

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

#26541 - Fix the type annotations and don't += a generator (it's confusing)

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

#26540 - Cache `shape` property on `Mesh`.

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

#26540 - Cache `shape` property on `Mesh`.

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

#26540 - Cache `shape` property on `Mesh`.

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

#26539 - Minor format cleanup

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

#26538 - Remove an unused import

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

#26538 - Remove an unused import

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

#26538 - Remove an unused import

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

GitHub / jax-ml/jax issues and pull requests