Skip to content

Pull requests: jax-ml/jax

New pull request New
397 Open 15,575 Closed
397 Open 15,575 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

[Mosaic GPU] Use absl::StrCat instead std::string::operator+
#24498 opened Oct 24, 2024 by copybara-service bot Loading…
@apaszke
fix the wrong output of pallas attention kernel when q_len!=kv_len
#24495 opened Oct 24, 2024 by zinccat Loading…
1
Fix float0 behavior inside shard_map transpose under scan.
#24493 opened Oct 24, 2024 by copybara-service bot Loading…
@pschuh
Remove obsolete implements() decorator & fix tests pull ready Ready for copybara import and testing
#24492 opened Oct 23, 2024 by jakevdp Loading…
1
@jakevdp
Perform searchsorted binary search using unsigned intermediate values. pull ready Ready for copybara import and testing
#24490 opened Oct 23, 2024 by hawkinsp Loading…
@hawkinsp
1
Improve docs for jnp.arctan2 documentation pull ready Ready for copybara import and testing
#24488 opened Oct 23, 2024 by jakevdp Loading…
@jakevdp
Add build commands to the new CI script for running Bazel GPU presubmits
#24485 opened Oct 23, 2024 by copybara-service bot Loading…
Add new CI scripts for running Bazel CPU presubmits
#24484 opened Oct 23, 2024 by copybara-service bot Loading…
Error on numpy array conversion of PRNG key array pull ready Ready for copybara import and testing
#24481 opened Oct 23, 2024 by jakevdp Loading…
@jakevdp
Run dot algorithm tests with PJRT plugin pull ready Ready for copybara import and testing
#24480 opened Oct 23, 2024 by dfm Loading…
@dfm
sharding cleanup: use inline checks for unimplemented and auto pull ready Ready for copybara import and testing
#24470 opened Oct 22, 2024 by jakevdp Loading…
@jakevdp
Add newline at the end of .bazelrc pull ready Ready for copybara import and testing
#24466 opened Oct 22, 2024 by andportnoy Loading…
[pallas:mosaic_gpu] Extracted a basic emit_pipeline API from the in kernel pipelining test
#24459 opened Oct 22, 2024 by copybara-service bot Loading…
@superbobry
[Mosaic] Add extension files for infer/apply vector layout.
#24444 opened Oct 22, 2024 by copybara-service bot Loading…
#sdy add shardy CPU config for all JAX tests, disabling any known failing test cases.
#24440 opened Oct 21, 2024 by copybara-service bot Loading…
Alternative abs() formula for sph_harm() with certain GPU/CUDA combinations
#24438 opened Oct 21, 2024 by yhtang Loading…
Update hermetic CUDA docs.
#24436 opened Oct 21, 2024 by copybara-service bot Loading…
Add support for layouts and other advanced features in ffi_call
#24433 opened Oct 21, 2024 by dfm Loading…
@dfm
3
Retry connection errors to combat flakiness in google metadata server during tpu initialization
#24432 opened Oct 21, 2024 by dlwh Loading…
1
Bump hypothesis from 6.102.4 to 6.115.3 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#24430 opened Oct 21, 2024 by dependabot bot Loading…
Bump fsspec from 2024.5.0 to 2024.10.0 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#24429 opened Oct 21, 2024 by dependabot bot Loading…
[NVIDIA] Adjust the atol/rtol of JAX SDPA API tests pull ready Ready for copybara import and testing
#24427 opened Oct 21, 2024 by kaixih Loading…
2
[pallas] Removed deprecated aliases to CostEstimate and run_scoped
#24419 opened Oct 21, 2024 by copybara-service bot Loading…
@superbobry
Add an example for executing serialized function in cpp
#24413 opened Oct 20, 2024 by fenghuizhang Loading…
8
Clean up conditional lowering of tan after JAX v0.4.34 release. pull ready Ready for copybara import and testing
#24405 opened Oct 19, 2024 by dfm Loading…
@dfm
ProTip! Exclude everything labeled bug with -label:bug.