jax-ml/jax issues and pull requests

#24065 - docs: update Custom_Operation_for_GPUs.md

Pull Request - State: closed - Opened by eltociear about 1 month ago
Labels: pull ready

#24064 - improve readthedocs behavior for jax.remat / jax.checkpoint

Pull Request - State: open - Opened by mattjj about 1 month ago
Labels: documentation

#24063 - [Mosaic TPU][NFC] Refactor tpu matmul rule in apply vector layout pass.

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

#24062 - improve jax.checkpoint / jax.remat prevent_cse docstring

Pull Request - State: open - Opened by mattjj about 1 month ago

#24061 - [Mosaic TPU] Change getLayout to force offset to 0 when inferring input has offset out of the first tile.

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

#24061 - [Mosaic TPU] Change getLayout to force offset to 0 when inferring input has offset out of the first tile.

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

#24059 - Edge case: Normal CDF incorrect when called with extremely small integer value if location is specified as integer array

Issue - State: open - Opened by jleugeri about 1 month ago - 1 comment
Labels: bug

#24058 - [Pallas] Add runtime assert via checkify.check. This check will halt the TPU if triggered, meaning that we would need to restart the program to recover.

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

#24057 - Allow simple host call that uses host tensor as parameter/result.

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

#24053 - Non-closure version of jax.vjp

Issue - State: open - Opened by luyug about 1 month ago - 3 comments
Labels: enhancement

#24051 - Regression in JAX FP8 matmul fusion

Issue - State: open - Opened by balancap about 1 month ago - 4 comments
Labels: bug

#24049 - Upgrade JAX jit Automatic Parallelization docs

Pull Request - State: open - Opened by 8bitmp3 about 1 month ago

#24047 - dot_product_attention has inconsistent dot types in backward pass

Issue - State: open - Opened by sbodenstein about 1 month ago - 6 comments
Labels: bug

#24044 - [Pallas/MGPU] Add support for tiled and swizzled loads/stores + support slices

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

#24039 - add checkify rule for remat

Pull Request - State: closed - Opened by mattjj about 1 month ago
Labels: pull ready

#24032 - Add timeout to requests calls

Pull Request - State: closed - Opened by pixeeai about 1 month ago - 5 comments
Labels: pull ready

#24031 - Fix vmap error message when args passed by keyword

Pull Request - State: closed - Opened by garymm about 1 month ago - 2 comments
Labels: bug, pull ready

#24027 - [Pallas TPU] Core dump when using `jnp.remainder`

Issue - State: closed - Opened by ayaka14732 about 1 month ago - 3 comments
Labels: bug, pallas

#24022 - [Mosaic] Parameterize the number of lanes and sublanes in TPU dialects.

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

#24017 - Bump hypothesis from 6.102.4 to 6.112.2

Pull Request - State: closed - Opened by dependabot[bot] about 2 months ago - 1 comment
Labels: dependencies, python

#24015 - Bump fsspec from 2024.5.0 to 2024.9.0

Pull Request - State: closed - Opened by dependabot[bot] about 2 months ago - 1 comment
Labels: dependencies, python

#24002 - [export] Enable more cross-platform lowering tests for GPU.

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

#23998 - Added `scatter_sub_p`

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

#23997 - [memories] Enable already-passing ComputeOffload tests on GPU

Pull Request - State: closed - Opened by gspschmid about 2 months ago - 4 comments

#23995 - Update docs for jnp.trapezoid

Pull Request - State: open - Opened by jakevdp about 2 months ago
Labels: documentation

#23994 - [Mosaic] Extend tpu matmulop to have dimension dims. Add support for batching and simple transposition.

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

#23993 - Automated Code Change

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

#23992 - [DOC] Incorrect interval in `jax.nn.log_softmax`

Issue - State: open - Opened by Qazalbash about 2 months ago
Labels: enhancement

#23989 - Mosaic failed to compile TPU kernel

Issue - State: open - Opened by jhkchan about 2 months ago - 2 comments
Labels: bug

#23988 - [Pallas TPU] Unrelated 'maximum recursion depth exceeded' error when using `.astype(jnp.int64)`

Issue - State: open - Opened by ayaka14732 about 2 months ago
Labels: bug, pallas

#23987 - [Mosaic TPU] Force offset to 0 when inferring input has offset out of the first tile.

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

#23986 - [Mosaic TPU] Canonicalize arith.select's condition to vector if other types are vector.

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

#23985 - Fix typo in Pallas TPU matmul doc. I think the logical layout of the input array is non-transposed, rather than transposed?

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

#23984 - jnp.mask_indices: add docs & tests

Pull Request - State: open - Opened by jakevdp about 2 months ago
Labels: documentation

#23983 - Bump test tolerance on FFT test that started failing in CI after an XLA change.

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

#23982 - Add support for FFI calls with side effects via ffi_call

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

#23981 - actions/checkout: pin specific hash

Pull Request - State: closed - Opened by jakevdp about 2 months ago - 1 comment
Labels: pull ready

#23980 - actions: remove use of ratchet tool

Pull Request - State: closed - Opened by jakevdp about 2 months ago
Labels: pull ready

#23979 - Fix test failure when shardy is not enabled.

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

#23978 - Include the sdy MLIR dialect in jaxlib.

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

#23977 - DOC: Improve docs for `jax.numpy`: `arcsinh`, `arccosh` and `arctanh`

Pull Request - State: closed - Opened by rajasekharporeddy about 2 months ago - 2 comments
Labels: documentation, pull ready

#23976 - Test PR.

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

#23974 - [Pallas/MGPU] Undo transforms on refs before giving them back to the users

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

#23973 - Fix tests that ask for an accelerator but don't use it.

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

#23972 - [Pallas TPU] `jnp.uint32` comparisons give the wrong result with large numbers

Issue - State: open - Opened by ayaka14732 about 2 months ago
Labels: bug, pallas

#23971 - Ahead-of-time compilation on GPU fails without a GPU device

Issue - State: open - Opened by jaro-sevcik about 2 months ago - 1 comment
Labels: bug

#23969 - [jax.distributed] Allow enabling grpc channel compression

Pull Request - State: open - Opened by gspschmid about 2 months ago

#23968 - Run Pallas Mosaic GPU tests on internal CI

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

#23967 - [Pallas TPU] Remove the skipTest for OpsExtraTest on TPU

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

#23966 - Removed unused imports in `jax.experimental.mosaic.gpu.core`

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

#23965 - Enable weight offloading tests that are supported on GPUs now

Pull Request - State: closed - Opened by zhenying-liu about 2 months ago
Labels: pull ready

#23964 - Update device_memory_profiling.md

Pull Request - State: open - Opened by garymm about 2 months ago

#23961 - Remove pylint sections from pyproject.toml (use ruff instead)

Pull Request - State: closed - Opened by apivovarov about 2 months ago - 2 comments
Labels: pull ready

#23960 - [Pallas] Simplify `OpsTest` by skipping 64-bit tests on 32-bit environments

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

#23959 - Add .pylintrc file

Pull Request - State: closed - Opened by apivovarov about 2 months ago - 3 comments

#23958 - jitted function gives incorrect result in jnp.concatenate

Issue - State: open - Opened by phdum-a about 2 months ago
Labels: bug

#23957 - Better documentation for jnp.choose

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

#23956 - Exposed sequential iteration index via `pl.program_id` in Pallas Mosaic GPU

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

#23955 - Fixed Pallas Mosaic GPU test following recent changes

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

#23954 - [XLA:Mosaic][Pallas] Enable vector.ExtractOp for non-zero indices.

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

#23953 - Replace disable_backends with enable_backends on jax_multiplatform_test.

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

#23951 - Add support for passing array attributes via ffi_call and an example demonstrating the different attribute APIs supported by the FFI

Pull Request - State: closed - Opened by dfm about 2 months ago - 6 comments
Labels: pull ready

GitHub / jax-ml/jax issues and pull requests