Ecosyste.ms: Issues
An open API service for providing issue and pull request metadata for open source projects.
GitHub / jax-ml/jax issues and pull requests
#24484 - Add new CI scripts for running Bazel CPU presubmits
Pull Request -
State: open - Opened by copybara-service[bot] 22 days ago
#24481 - Error on numpy array conversion of PRNG key array
Pull Request -
State: open - Opened by jakevdp 22 days ago
Labels: pull ready
#24480 - Run dot algorithm tests with PJRT plugin
Pull Request -
State: open - Opened by dfm 22 days ago
Labels: pull ready
#24478 - Attribute Error
Issue -
State: closed - Opened by Opiyo-Inno 22 days ago
- 17 comments
Labels: bug
#24473 - Merge v0.4.35 release branch
Pull Request -
State: open - Opened by hawkinsp 22 days ago
Labels: pull ready
#24470 - sharding cleanup: use inline checks for unimplemented and auto
Pull Request -
State: closed - Opened by jakevdp 23 days ago
Labels: pull ready
#24469 - Better documentation for jnp.indices
Pull Request -
State: open - Opened by jakevdp 23 days ago
Labels: documentation, pull ready
#24468 - Use libtpu releases rather than libtpu-nightly for jax[tpu].
Pull Request -
State: closed - Opened by copybara-service[bot] 23 days ago
#24467 - Add back the import of `jtu` in `flash_attention.py`
Pull Request -
State: closed - Opened by andportnoy 23 days ago
- 1 comment
Labels: pull ready
#24466 - Add newline at the end of `.bazelrc`
Pull Request -
State: open - Opened by andportnoy 23 days ago
Labels: pull ready
#24465 - Fix mypy error
Pull Request -
State: closed - Opened by jakevdp 23 days ago
Labels: pull ready
#24464 - [Pallas TPU] Unable to create a boolean array inside kernel due to Mosaic relayout error
Issue -
State: open - Opened by ayaka14732 23 days ago
Labels: bug, pallas
#24463 - Require ml_dtypes >= 0.4.0
Pull Request -
State: closed - Opened by jakevdp 23 days ago
Labels: pull ready
#24462 - [sharding_in_types][Take 2] Add `out_type` argument to `einsum` and `dot_general` to allow specifying for the output type. Right now, it only accept a `NamedSharding` but in the future we can allow a polymorphic type of: `jax.ShapeDtypeStruct | Sharding | Layout`.
Pull Request -
State: closed - Opened by copybara-service[bot] 23 days ago
#24461 - Should ml-dtypes be in the requirements?
Issue -
State: closed - Opened by cottrell 23 days ago
- 4 comments
Labels: bug
#24460 - Full coverage for `__jax_array__` protocol
Issue -
State: open - Opened by maxencefaldor 23 days ago
- 1 comment
Labels: bug
#24459 - [pallas:mosaic_gpu] Extracted a basic `emit_pipeline` API from the in kernel pipelining test
Pull Request -
State: closed - Opened by copybara-service[bot] 23 days ago
#24458 - [ROCM] Sharding leads to HIP errors
Issue -
State: open - Opened by PhilipVinc 23 days ago
Labels: bug
#24457 - [Pallas:MGPU] Treat each warpgroup as a single logical thread.
Pull Request -
State: open - Opened by copybara-service[bot] 23 days ago
#24456 - [Pallas] Allow core_map's mesh to discharge backend specific effects
Pull Request -
State: open - Opened by copybara-service[bot] 23 days ago
#24455 - [Pallas:MGPU] Allow initializing accumulators with values in registers
Pull Request -
State: open - Opened by copybara-service[bot] 23 days ago
#24454 - Inconsistency in using == operator on a jax array and a list
Issue -
State: open - Opened by gokul-uf 23 days ago
- 1 comment
Labels: bug
#24452 - Better docs for jnp.insert
Pull Request -
State: closed - Opened by jakevdp 23 days ago
Labels: documentation, pull ready
#24451 - Add an extension mechanism to run_state that allows:
Pull Request -
State: open - Opened by copybara-service[bot] 23 days ago
#24450 - [Pallas:MGPU] Add support for grid dims in GPUMesh
Pull Request -
State: closed - Opened by copybara-service[bot] 23 days ago
#24449 - [Mosaic GPU] Query SM and PTX ISA dynamically using driver and LLVM
Pull Request -
State: open - Opened by copybara-service[bot] 23 days ago
#24447 - Reverts ebb75db8a523150c48376d15391f84380a2bb110
Pull Request -
State: closed - Opened by copybara-service[bot] 23 days ago
#24446 - [docs] Added two new APIs to the export API docs
Pull Request -
State: open - Opened by gnecula 23 days ago
Labels: pull ready
#24445 - Raise error when out_type is specified but sharding_in_types is False
Pull Request -
State: closed - Opened by copybara-service[bot] 23 days ago
#24444 - [Mosaic] Add extension files for infer/apply vector layout.
Pull Request -
State: open - Opened by copybara-service[bot] 23 days ago
#24443 - [sharding_in_types] Add `out_type` argument to `einsum` and `dot_general` to allow specifying for the output type. Right now, it only accept a `NamedSharding` but in the future we can allow a polymorphic type of: `jax.ShapeDtypeStruct | Sharding | Layout`.
Pull Request -
State: closed - Opened by copybara-service[bot] 23 days ago
#24442 - Better documentation for jnp.lexsort
Pull Request -
State: closed - Opened by jakevdp 23 days ago
Labels: documentation, pull ready
#24441 - [pallas:mosaic_gpu] An extremely specific heuristic to allow swiglu.
Pull Request -
State: open - Opened by copybara-service[bot] 23 days ago
#24440 - #sdy add shardy CPU config for all JAX tests, disabling any known failing test cases.
Pull Request -
State: open - Opened by copybara-service[bot] 23 days ago
#24439 - Fix documentation.
Pull Request -
State: closed - Opened by copybara-service[bot] 23 days ago
#24438 - Alternative abs() formula for `sph_harm()` with certain GPU/CUDA combinations
Pull Request -
State: open - Opened by yhtang 23 days ago
#24437 - Fix jnp.diff for boolean inputs
Pull Request -
State: closed - Opened by jakevdp 23 days ago
Labels: pull ready
#24436 - Update hermetic CUDA docs.
Pull Request -
State: closed - Opened by copybara-service[bot] 23 days ago
#24435 - pmap of pjit
Issue -
State: open - Opened by luyug 24 days ago
- 4 comments
Labels: enhancement
#24434 - Update Pallas distributed tutorials with jax.make_mesh
Pull Request -
State: closed - Opened by justinjfu 24 days ago
Labels: pull ready
#24433 - Add support for layouts and other advanced features in ffi_call
Pull Request -
State: closed - Opened by dfm 24 days ago
- 3 comments
Labels: pull ready
#24432 - Retry connection errors to combat flakiness in google metadata server during tpu initialization
Pull Request -
State: open - Opened by dlwh 24 days ago
- 1 comment
#24430 - Bump hypothesis from 6.102.4 to 6.115.3
Pull Request -
State: open - Opened by dependabot[bot] 24 days ago
Labels: dependencies, python
#24429 - Bump fsspec from 2024.5.0 to 2024.10.0
Pull Request -
State: open - Opened by dependabot[bot] 24 days ago
Labels: dependencies, python
#24428 - Add test for jax2tf conversion of dot general with algorithm.
Pull Request -
State: open - Opened by copybara-service[bot] 24 days ago
#24427 - [NVIDIA] Adjust the atol/rtol of JAX SDPA API tests
Pull Request -
State: open - Opened by kaixih 24 days ago
- 1 comment
Labels: pull ready
#24425 - Rename vmap methods for callbacks
Pull Request -
State: closed - Opened by dfm 24 days ago
Labels: pull ready
#24424 - FIx jax2tf breakge of iota
Pull Request -
State: closed - Opened by copybara-service[bot] 24 days ago
#24422 - [Pallas:MGPU] Allow allocating transformed refs in run_scoped
Pull Request -
State: closed - Opened by copybara-service[bot] 24 days ago
#24421 - Better docs for jnp.cross
Pull Request -
State: closed - Opened by jakevdp 24 days ago
Labels: documentation, pull ready
#24420 - Bumped mypy and ruff versions used by pre-commit
Pull Request -
State: closed - Opened by superbobry 24 days ago
Labels: pull ready
#24419 - [pallas] Removed deprecated aliases to `CostEstimate` and `run_scoped`
Pull Request -
State: closed - Opened by copybara-service[bot] 24 days ago
#24418 - shard_map - NotImplementedError: No replication rule for cond
Issue -
State: open - Opened by giovannicemin 24 days ago
- 1 comment
Labels: bug
#24417 - Update warning message for jit-of-pmap
Pull Request -
State: closed - Opened by rajasekharporeddy 24 days ago
Labels: pull ready
#24415 - Not implemented: Non-trivial layouts unsupported
Issue -
State: open - Opened by vanbasten23 24 days ago
- 4 comments
Labels: bug, pallas
#24412 - Better docs for jnp.fromfunction
Pull Request -
State: closed - Opened by jakevdp 25 days ago
Labels: documentation, pull ready
#24411 - JAX code is extremely slow on GPUs
Issue -
State: open - Opened by AakashKumarNain 25 days ago
- 23 comments
Labels: bug, performance
#24408 - docs: supported platform info incorrect for Mac
Issue -
State: closed - Opened by garymm 25 days ago
Labels: bug
#24406 - vmap gives misleading error message about inconsistent sizes when passing arg by keyword
Issue -
State: closed - Opened by garymm 25 days ago
Labels: bug
#24405 - Clean up conditional lowering of tan after JAX v0.4.34 release.
Pull Request -
State: open - Opened by dfm 25 days ago
Labels: pull ready
#24404 - INTERNAL: Mosaic failed to compile TPU kernel: unsupported shape cast
Issue -
State: open - Opened by vanbasten23 26 days ago
- 4 comments
Labels: bug, pallas
#24403 - Better documentation for jnp.load
Pull Request -
State: closed - Opened by jakevdp 26 days ago
Labels: documentation, pull ready
#24400 - Make jnp.subtract a ufunc
Pull Request -
State: closed - Opened by jakevdp 26 days ago
Labels: documentation, pull ready
#24399 - jax.distributed.initialize() crash
Issue -
State: open - Opened by demon2036 26 days ago
- 4 comments
Labels: bug
#24398 - jax.jit changes the key order of returned dictionaries
Issue -
State: open - Opened by carlosgmartin 26 days ago
- 6 comments
Labels: bug
#24393 - Fix env check order by splitting TPU envs into separate files.
Pull Request -
State: closed - Opened by copybara-service[bot] 27 days ago
#24392 - Respect dot algorithm spec on TPU backends.
Pull Request -
State: closed - Opened by copybara-service[bot] 27 days ago
#24391 - Document jax config to disable remat HLO pass
Pull Request -
State: closed - Opened by keshavb96 27 days ago
- 2 comments
Labels: pull ready
#24390 - Allowing HLO constants inside custom partition
Issue -
State: open - Opened by ASKabalan 27 days ago
- 4 comments
Labels: enhancement
#24389 - [Pallas:MGPU] Add support for passing in WGMMA lhs from registers
Pull Request -
State: closed - Opened by copybara-service[bot] 27 days ago
#24383 - [Pallas:MGPU] Allow indexing to appear anywhere in the list of transforms
Pull Request -
State: closed - Opened by copybara-service[bot] 27 days ago
#24378 - Register acos primitive to lower to CHLO acos.
Pull Request -
State: closed - Opened by copybara-service[bot] 27 days ago
#24370 - Expand the `ffi_call` API to support more use cases
Pull Request -
State: closed - Opened by dfm 28 days ago
Labels: pull ready
#24368 - [pallas] `input_output_aliases` now only include refs which have been written to
Pull Request -
State: open - Opened by copybara-service[bot] 28 days ago
#24356 - [Pallas TPU] Fix lowering for `jnp.remainder`
Pull Request -
State: closed - Opened by copybara-service[bot] 28 days ago
#24355 - [export] Fix github links in the export documentation
Pull Request -
State: closed - Opened by gnecula 28 days ago
Labels: pull ready
#24354 - [Mosaic TPU] Support dynamic DMA and ref slice on the 2nd minor when memref is untiled
Pull Request -
State: closed - Opened by copybara-service[bot] 28 days ago
#24352 - [NVIDIA] Support custom dtype convert in jax.nn.dot_product_attention
Pull Request -
State: open - Opened by kaixih 28 days ago
- 2 comments
#24349 - CPU profiling (not tracing)
Issue -
State: open - Opened by joaospinto 28 days ago
- 5 comments
Labels: enhancement
#24346 - Move converter related tflite functions to tensorflow/lite repo
Pull Request -
State: closed - Opened by copybara-service[bot] 29 days ago
#24342 - [export] Add support for serialization for some custom PyTree nodes
Pull Request -
State: closed - Opened by gnecula 29 days ago
Labels: pull ready
#24310 - [Pallas TPU] Add lowerings for bf16 `jnp.ceil` and `jnp.floor` in TPU v6+
Pull Request -
State: closed - Opened by copybara-service[bot] 30 days ago
#24304 - Rollback v5e device change
Pull Request -
State: closed - Opened by copybara-service[bot] about 1 month ago
#24302 - resolve extended dtypes in a jaxpr pass
Pull Request -
State: open - Opened by froystig about 1 month ago
Labels: pull ready
#24298 - Move out mnist py/Jax tensorflow lite tests to tensorflow lite repo
Pull Request -
State: closed - Opened by copybara-service[bot] about 1 month ago
#24296 - Remove silent data corruption runtime flags from persistent cache key.
Pull Request -
State: closed - Opened by copybara-service[bot] about 1 month ago
#24293 - Update landing page styling
Pull Request -
State: open - Opened by barnesjoseph about 1 month ago
#24291 - Bump hypothesis from 6.102.4 to 6.115.0
Pull Request -
State: closed - Opened by dependabot[bot] about 1 month ago
- 1 comment
Labels: dependencies, python
#24277 - [pallas] Support for setting explicit backends to pallas_call.
Pull Request -
State: closed - Opened by copybara-service[bot] about 1 month ago
#24268 - Fix repr of sharding in aval when a dimension is sharded on multiple mesh axes
Pull Request -
State: closed - Opened by copybara-service[bot] about 1 month ago
#24267 - [sharding_in_types] Add constraints during lowering for dot_general and reduce_sum so that we can enforce the sharding we choose during tracing
Pull Request -
State: closed - Opened by copybara-service[bot] about 1 month ago
#24266 - [sharding_in_types] Add sharding rule for reduce sum which is just drop the specs for the axis we are reducing over
Pull Request -
State: closed - Opened by copybara-service[bot] about 1 month ago
#24264 - [ROCm] apt update
Pull Request -
State: open - Opened by Ruturaj4 about 1 month ago
Labels: pull ready
#24263 - Better repr of aval when shardings are present
Pull Request -
State: closed - Opened by copybara-service[bot] about 1 month ago
#24261 - [sharding_in_types] Add `dot_general` sharding rule. We only handle the simple cases and rely on xla to insert the collectives.
Pull Request -
State: closed - Opened by copybara-service[bot] about 1 month ago
#24260 - [XLA:Collective] Expose a factory for constructing HLOSharding with explicit device ordering.
Pull Request -
State: open - Opened by copybara-service[bot] about 1 month ago
#24259 - `lax.cond` is much slower than calling `true_fn` itself when `cond` is always true.
Issue -
State: open - Opened by hanzhi713 about 1 month ago
Labels: enhancement
#24258 - Remove checks for jaxlib v0.4.33 in tests
Pull Request -
State: open - Opened by dfm about 1 month ago
#24257 - Don't add `mhlo.layout_mode = "default"` since that is the default even in PJRT and will help reduce cruft in the IR
Pull Request -
State: closed - Opened by copybara-service[bot] about 1 month ago
#24256 - Remove temporary aliases from .bazelrc
Pull Request -
State: closed - Opened by copybara-service[bot] about 1 month ago