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
#24756 - pure_callback using a threadpool to accelerate vmap
Issue -
State: open - Opened by Joshuaalbert 7 days ago
- 2 comments
Labels: enhancement
#24755 - [mosaic_gpu] The profiler now uses FFI calls for creating events and computing elapsed time
Pull Request -
State: closed - Opened by copybara-service[bot] 7 days ago
#24753 - Facing Scaling issue on cpu (arm and x86).
Issue -
State: open - Opened by choudhary-devang 7 days ago
Labels: bug
#24752 - Propagate CopySemantics from python to C++ transfer APIs so that device_put works correctly in presence of copy/donate options that user specified.
Pull Request -
State: closed - Opened by copybara-service[bot] 7 days ago
#24751 - Allow setting default_device with platform names.
Pull Request -
State: open - Opened by Stella-S-Yan 7 days ago
- 1 comment
Labels: pull ready
#24750 - disable remat HLO pass by default
Pull Request -
State: open - Opened by keshavb96 7 days ago
- 1 comment
#24749 - cudaErrorSymbolNotFound : named symbol not found
Issue -
State: open - Opened by amanjitsk 7 days ago
- 3 comments
Labels: bug
#24748 - jnp.reshape: raise TypeError when specifying newshape
Pull Request -
State: open - Opened by jakevdp 7 days ago
#24746 - [Pallas] Lower `lax.exp2` to MLIR `math.exp2` directly
Pull Request -
State: open - Opened by copybara-service[bot] 8 days ago
#24741 - [Mosaic] Fix canonicalize_extract op name.
Pull Request -
State: closed - Opened by copybara-service[bot] 8 days ago
#24740 - TracerBoolConversionError
Issue -
State: closed - Opened by hangita101 8 days ago
- 4 comments
Labels: question
#24738 - debug export
Pull Request -
State: closed - Opened by copybara-service[bot] 8 days ago
#24737 - `scipy.sparse.csgraph.connected_components` implementation
Issue -
State: open - Opened by vboussange 8 days ago
- 2 comments
Labels: enhancement
#24735 - [Mosaic] Fix handling of i1 splat constants
Pull Request -
State: closed - Opened by copybara-service[bot] 8 days ago
#24734 - More thorough propagation of host linear layout
Pull Request -
State: open - Opened by copybara-service[bot] 8 days ago
#24730 - InconclusiveDimensionOperation: Symbolic dimension comparison 'b' < '2147483647' is inconclusive.
Issue -
State: closed - Opened by njzjz 8 days ago
- 3 comments
Labels: bug
#24724 - [SDY] add JAX lowering to Shardy `ShardingGroupOp` for shard_alike.
Pull Request -
State: closed - Opened by copybara-service[bot] 9 days ago
#24721 - Expose the HloOpcode::kPartitionId in Jax.
Issue -
State: closed - Opened by ptoulme-aws 9 days ago
- 7 comments
Labels: enhancement
#24718 - [Mosaic GPU] Implement tiled and swizzled transfers for tiled layouts
Pull Request -
State: closed - Opened by copybara-service[bot] 9 days ago
#24714 - [Mosaic GPU] Introduce a more flexible layout system
Pull Request -
State: closed - Opened by copybara-service[bot] 9 days ago
#24706 - [XLA] Make our LLVM usage more googley
Pull Request -
State: open - Opened by copybara-service[bot] 9 days ago
#24705 - RFC: specify jit static args via Static annotation
Pull Request -
State: open - Opened by jakevdp 9 days ago
- 3 comments
Labels: enhancement
#24700 - register_dataclass: allow marking static fields via dataclass.field
Pull Request -
State: closed - Opened by jakevdp 9 days ago
Labels: pull ready
#24696 - Bump hypothesis from 6.102.4 to 6.116.0
Pull Request -
State: closed - Opened by dependabot[bot] 10 days ago
- 1 comment
Labels: dependencies, python
#24695 - Re-factor build CLI to a subcommand based approach
Pull Request -
State: open - Opened by copybara-service[bot] 10 days ago
#24682 - Disable 64 bit codegen to HLO for Neuron backend
Issue -
State: open - Opened by ptoulme-aws 10 days ago
- 7 comments
Labels: enhancement
#24674 - [mosaic_gpu] Scalar arguments to kernels.
Pull Request -
State: closed - Opened by copybara-service[bot] 12 days ago
#24670 - [Mosaic GPU] Implement the skeleton of a lowering pass for the Mosaic GPU dialect.
Pull Request -
State: closed - Opened by copybara-service[bot] 13 days ago
#24664 - RFC: add pytree-compatible dataclass decorator
Pull Request -
State: open - Opened by jakevdp 13 days ago
- 3 comments
#24663 - Add a GPU implementation of `lax.linalg.eig`.
Pull Request -
State: open - Opened by copybara-service[bot] 13 days ago
- 1 comment
#24657 - Fix argmin docstring to not say "maximum"
Pull Request -
State: closed - Opened by dymil 13 days ago
Labels: pull ready
#24655 - Automated Code Change
Pull Request -
State: open - Opened by copybara-service[bot] 13 days ago
#24648 - `jax.lax.dot_general`: getting `failed to legalize operation 'vhlo.dot_general_v2' that was explicitly marked illegal`
Issue -
State: closed - Opened by andportnoy 13 days ago
- 4 comments
Labels: bug
#24645 - [Mosaic GPU] Improve correctness of benchmarking scripts
Pull Request -
State: closed - Opened by andportnoy 14 days ago
Labels: pull ready
#24630 - [Pallas:MGPU] Add support for multiple heads in attention
Pull Request -
State: closed - Opened by copybara-service[bot] 15 days ago
#24629 - [Mosaic GPU] Fix the ordering of transforms in async_copy
Pull Request -
State: closed - Opened by copybara-service[bot] 15 days ago
#24602 - GPU Pallas `decode_attention` improvements
Pull Request -
State: closed - Opened by rdyro 15 days ago
Labels: pull ready
#24594 - Refactor JAX build wheel rule and add wheel_library targets.
Pull Request -
State: open - Opened by copybara-service[bot] 16 days ago
#24590 - [pallas:mosaic_gpu] Ported two pipelining optimizations to `emit_pipeline`
Pull Request -
State: open - Opened by copybara-service[bot] 16 days ago
#24561 - Plumb dot dimension numbers into TPU matmul op.
Pull Request -
State: closed - Opened by copybara-service[bot] 17 days ago
#24560 - Move utility functions in build.py to utils.py
Pull Request -
State: open - Opened by copybara-service[bot] 17 days ago
#24559 - [Mosaic GPU] Add Python bindings for the Mosaic GPU MLIR dialect.
Pull Request -
State: open - Opened by copybara-service[bot] 17 days ago
#24558 - Raise a better error message if `None` is passed to with_sharding_constraint.
Pull Request -
State: open - Opened by copybara-service[bot] 17 days ago
#24557 - Add base case to _device_id_to_logical.
Pull Request -
State: closed - Opened by copybara-service[bot] 17 days ago
#24556 - Fix breaking PGLE test-cases
Pull Request -
State: closed - Opened by copybara-service[bot] 17 days ago
#24555 - Skip testVectorizedDeprecation on Python 3.13 to unblock the CI
Pull Request -
State: closed - Opened by copybara-service[bot] 17 days ago
#24554 - Remove int4 from jtu.dtypes.all_dtypes.
Pull Request -
State: open - Opened by jakevdp 17 days ago
#24553 - [Mosaic GPU] Don't assume register vectors are exactly 2 in length
Pull Request -
State: closed - Opened by copybara-service[bot] 17 days ago
#24552 - [Pallas:MGPU] Fix flaky debug_print tests
Pull Request -
State: closed - Opened by copybara-service[bot] 17 days ago
#24551 - Extend PR23304 to include jnp.maximum and jnp.minimum
Pull Request -
State: open - Opened by copybara-service[bot] 17 days ago
#24549 - MaxText Getting Started fails
Issue -
State: open - Opened by learning-to-play 17 days ago
- 1 comment
Labels: bug
#24548 - segmentation fault from a simple test on GCP vm
Issue -
State: open - Opened by jiyuuchc 17 days ago
- 2 comments
Labels: bug
#24547 - New function to allow easy paralelization over vmap()
Issue -
State: open - Opened by MRiabov 18 days ago
- 4 comments
Labels: enhancement
#24546 - [pallas:mosaic_gpu] `emit_pipeline` now maintains the grid indices
Pull Request -
State: closed - Opened by copybara-service[bot] 18 days ago
#24545 - fix jax.custom_gradient to allow closing over non-autodiff tracers
Pull Request -
State: open - Opened by mattjj 18 days ago
Labels: pull ready
#24544 - ⚠️ Nightly upstream-dev CI failed ⚠️
Issue -
State: open - Opened by github-actions[bot] 19 days ago
- 2 comments
Labels: CI
#24543 - Update pre-commit workflow to cache on jax version
Pull Request -
State: open - Opened by dfm 19 days ago
Labels: pull ready
#24542 - Fix lint at head
Pull Request -
State: closed - Opened by dfm 19 days ago
Labels: pull ready
#24541 - return in finally swallows exceptions
Issue -
State: open - Opened by iritkatriel 19 days ago
- 1 comment
Labels: bug
#24540 - Remove the `initial` argument to `jax.nn.softmax` and `jax.nn.log_softmax`.
Pull Request -
State: open - Opened by copybara-service[bot] 19 days ago
#24539 - [Pallas][Flash Attention] Cost analysis in Flash Attention kernel causing compilation OOM
Issue -
State: open - Opened by lsy323 19 days ago
Labels: bug
#24538 - Add new jnp.cumulative_prod function.
Pull Request -
State: closed - Opened by jakevdp 19 days ago
Labels: pull ready
#24537 - Add some missing jax.numpy documentation
Pull Request -
State: closed - Opened by jakevdp 19 days ago
Labels: documentation, pull ready
#24536 - Limit self-hosted jobs to JAX main repo
Pull Request -
State: open - Opened by MichaelHudgins 20 days ago
#24535 - #jax Adds a missing comma to Pallas Quickstart
Pull Request -
State: closed - Opened by copybara-service[bot] 20 days ago
#24534 - Update jaxlib GPU FFI handlers to use templated `typed_data` method rather than explicit casting.
Pull Request -
State: open - Opened by copybara-service[bot] 20 days ago
#24533 - [Mosaic] Use max sublane offset per shuffled load to decide whether to avoid bank conflict.
Pull Request -
State: closed - Opened by copybara-service[bot] 20 days ago
#24531 - [mosaic] Directly build IR in _device_id_to_logical, rather than using lower_fun.
Pull Request -
State: closed - Opened by copybara-service[bot] 20 days ago
#24530 - [mosaic_gpu] `mgpu.FragmentedArray` now supports `//`
Pull Request -
State: closed - Opened by copybara-service[bot] 20 days ago
#24529 - [Pallas:MGPU] Add FlashAttention3 as an example
Pull Request -
State: open - Opened by copybara-service[bot] 20 days ago
#24528 - [Pallas:MGPU] Use shfl.sync after computing the warpgroup index
Pull Request -
State: open - Opened by copybara-service[bot] 20 days ago
#24527 - Fix a use-after-free bug in third_party/py/jax/jaxlib/mosaic/dialect/gpu/mosaic_gpu_test.cc
Pull Request -
State: closed - Opened by copybara-service[bot] 20 days ago
#24526 - [Pallas:MGPU] Wire up the Mosaic GPU profiler into Pallas
Pull Request -
State: closed - Opened by copybara-service[bot] 20 days ago
#24525 - Remove duplicate entry in CHANGELOG.
Pull Request -
State: closed - Opened by gnecula 20 days ago
Labels: pull ready
#24524 - Re-enable jax2tf test for dot algorithm with stricter TF version check
Pull Request -
State: closed - Opened by dfm 20 days ago
Labels: pull ready
#24523 - [ROCm] ci build fixes
Pull Request -
State: closed - Opened by Ruturaj4 20 days ago
- 1 comment
Labels: pull ready
#24522 - Disable jax2tf test recently added in cl/688976685.
Pull Request -
State: closed - Opened by copybara-service[bot] 20 days ago
#24521 - INVALID_ARGUMENT when copying array in pinned_host memory with device_put(..., may_alias=False)
Issue -
State: closed - Opened by ojw28 20 days ago
Labels: bug
#24520 - Add sharding rules to some more primitives so that backward pass of minformer passes. There are a couple of changes here:
Pull Request -
State: closed - Opened by copybara-service[bot] 20 days ago
#24519 - Add GC support for PyTreeRegistry.
Pull Request -
State: open - Opened by copybara-service[bot] 20 days ago
#24518 - Support None in PmapSharding as a replacement for device_put_replicated.
Pull Request -
State: closed - Opened by copybara-service[bot] 20 days ago
#24517 - [mosaic] Use .clone() to duplicate a module, rather than printing and parsing it.
Pull Request -
State: closed - Opened by copybara-service[bot] 20 days ago
#24516 - improve checkpoint / remat concreteness error with static_argnums
Pull Request -
State: open - Opened by mattjj 20 days ago
Labels: pull ready
#24515 - Add composite test to jax.
Pull Request -
State: open - Opened by copybara-service[bot] 20 days ago
#24514 - Clear extra_jit_context when exiting.
Pull Request -
State: closed - Opened by copybara-service[bot] 20 days ago
#24513 - [Mosaic TPU] Remove extra check
Pull Request -
State: closed - Opened by copybara-service[bot] 21 days ago
#24512 - Finalize deprecation of xb, xc, & xe symbols in jax.interpreters.xla
Pull Request -
State: closed - Opened by copybara-service[bot] 21 days ago
#24511 - improve concreteness error message in remat
Pull Request -
State: closed - Opened by mattjj 21 days ago
Labels: pull ready
#24510 - Allow dot algorithms in default_matmul_precision config
Pull Request -
State: open - Opened by dfm 21 days ago
#24509 - hlo_helpers.dtype_to_ir_type is not complete comparing with jax.interpreters.mlir.dtype_to_ir_type
Issue -
State: open - Opened by aws-qieqingy 21 days ago
- 7 comments
Labels: enhancement
#24507 - Finish 0.4.35 release by removing dead code
Pull Request -
State: closed - Opened by copybara-service[bot] 21 days ago
#24505 - [jax2tf] Fixes a bad interaction between jax2tf.convert, TF, and call_tf
Pull Request -
State: closed - Opened by gnecula 21 days ago
- 1 comment
Labels: pull ready
#24504 - [Pallas TPU] Add lowerings for scalar `absf` and `rsqrt`
Pull Request -
State: closed - Opened by copybara-service[bot] 21 days ago
#24500 - [shape_poly] Expands support for random.choice
Pull Request -
State: closed - Opened by gnecula 21 days ago
Labels: pull ready
#24499 - reduce_sum support for int4, int2 and uint4, uint2
Pull Request -
State: open - Opened by EdouardYvinec 21 days ago
- 2 comments
#24498 - [Mosaic GPU] Use absl::StrCat instead std::string::operator+
Pull Request -
State: closed - Opened by copybara-service[bot] 21 days ago
#24495 - fix the wrong output of pallas attention kernel when q_len!=kv_len
Pull Request -
State: closed - Opened by zinccat 21 days ago
Labels: pull ready
#24493 - Fix float0 behavior inside shard_map transpose under scan.
Pull Request -
State: closed - Opened by copybara-service[bot] 21 days ago
#24491 - Pallas Attention Kernel Wrong results when length of query != length of KV
Issue -
State: closed - Opened by zinccat 21 days ago
- 2 comments
Labels: bug
#24485 - Add build commands to the new CI script for running Bazel GPU presubmits
Pull Request -
State: open - Opened by copybara-service[bot] 22 days ago