Skip to content

Pull requests: jax-ml/jax

New pull request New
423 Open 16,971 Closed
423 Open 16,971 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

Check "jax_rocm_visible_devices" at client creation.
#26327 opened Feb 5, 2025 by ksebaz Loading…
1
Compute buffer aliasing on a per buffer basis.
#26323 opened Feb 5, 2025 by copybara-service bot Loading…
@pschuh
Lower JAX_NUM_GENERATED_CASES to 20.
#26322 opened Feb 5, 2025 by copybara-service bot Loading…
@nitins17
[JAX] Generate more readable error for failed device deserialization in colocated Python
#26321 opened Feb 4, 2025 by copybara-service bot Loading…
Run Bazel CPU/CUDA presubmit jobs tests on pushes to main/release branches
#26320 opened Feb 4, 2025 by copybara-service bot Loading…
@nitins17
Return arrays from ArrayImpl._check_and_rearrange.
#26316 opened Feb 4, 2025 by copybara-service bot Loading…
@emilyfertig
#sdy Fix the format for the external link to jax-shardy-guide colab.
#26315 opened Feb 4, 2025 by copybara-service bot Loading…
Test ci changes
#26314 opened Feb 4, 2025 by copybara-service bot Loading…
[better_errors] Add debug info to the Jaxprs formed for AD (step 2) pull ready Ready for copybara import and testing
#26313 opened Feb 4, 2025 by gnecula Loading…
@gnecula
[pallas:mgpu] Cast all indices to i32 during lowering.
#26311 opened Feb 4, 2025 by copybara-service bot Loading…
[Mosaic GPU] Infer whether A/B are row- or column-major from strides
#26306 opened Feb 4, 2025 by copybara-service bot Loading…
@apaszke
[pallas_mgpu] Fix emit_pipeline_with_wgmma test and update rtol.
#26301 opened Feb 4, 2025 by copybara-service bot Loading…
@petebu
[JAX] Add a test verifying the behavior of module-level state accessed by colocated Python
#26293 opened Feb 4, 2025 by copybara-service bot Loading…
Finalize some deprecations from jax.lib.xla_client
#26292 opened Feb 4, 2025 by jakevdp • Draft
@jakevdp
3
Extend random.orthogonal to semi-orthogonal matrices and simplify initializers.orthogonal pull ready Ready for copybara import and testing
#26291 opened Feb 4, 2025 by carlosgmartin Loading…
@jakevdp
8
[JAX][DOC] memory_kind, with_memory_kind and out_shardings
#26285 opened Feb 3, 2025 by zhenying-liu Loading…
@yashk2810
Proof-of-concept: Use XLA's intra_op_threadpool to batch parallelize JAX's LAPACK calls on CPU.
#26280 opened Feb 3, 2025 by copybara-service bot Loading…
@dfm
[CI] Add upload to resultstore to tsan bazel test run CI
#26279 opened Feb 3, 2025 by MichaelHudgins Loading…
1
[pallas:mgpu] Change FA3 kernel bc lax.div doesn't like mixed types anymore.
#26268 opened Feb 3, 2025 by copybara-service bot Loading…
Default JAX_CPU_COLLECTIVES_IMPLEMENTATION to 'gloo'.
#26264 opened Feb 2, 2025 by copybara-service bot Loading…
@skye
feat(gh-13291): Add exponential distribution functions: cdf, logcdf, sf, logsf, and ppf pull ready Ready for copybara import and testing
#26259 opened Feb 1, 2025 by Qazalbash Loading…
@jakevdp
6
Add check to lax.composite to prevent DynamicJaxprTracer type errors.
#26257 opened Feb 1, 2025 by copybara-service bot Loading…
Add num_updates dimension to input/output_offset and send/recv_sizes of ragged_all_to_all.
#26256 opened Feb 1, 2025 by copybara-service bot Loading…
Return arrays from ArrayImpl._check_and_rearrange. Build IFRT shardings with both addressable and non-addressable devices, instead of only addressable devices.
#26253 opened Jan 31, 2025 by copybara-service bot Loading…
@emilyfertig
Replace Python 3.12 with Python 3.13 and add Python 3.10 to the matrix
#26250 opened Jan 31, 2025 by copybara-service bot Loading…
@nitins17
ProTip! Filter pull requests by the default branch with base:main.