Skip to content

Pull requests: jax-ml/jax

New pull request New
424 Open 16,923 Closed
424 Open 16,923 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

Default JAX_CPU_COLLECTIVES_IMPLEMENTATION to 'gloo'.
#26264 opened Feb 2, 2025 by copybara-service bot Loading…
@skye
[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo pull ready Ready for copybara import and testing
#26262 opened Feb 1, 2025 by gnecula Loading…
@gnecula
Fix folder problem in TSAN numpy step
#26261 opened Feb 1, 2025 by vfdev-5 Loading…
Reverts 3f737588fd49fcac36d6bd00cfdabae07df1fbb2
#26260 opened Feb 1, 2025 by copybara-service bot Loading…
@hawkinsp
feat(gh-13291): Add exponential distribution functions: cdf, logcdf, sf, logsf, and ppf
#26259 opened Feb 1, 2025 by Qazalbash Loading…
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…
Propagate source ranges in location.
#26254 opened Jan 31, 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
Verify GKE cluster detection happens before generic k8s
#26251 opened Jan 31, 2025 by samos123 Loading…
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
jax.lax: improve docs for floor, ceil, round. documentation pull ready Ready for copybara import and testing
#26246 opened Jan 31, 2025 by jakevdp Loading…
@jakevdp
1
Make jnp.asarray lower to asarray_p in the simplest cases.
#26244 opened Jan 31, 2025 by jakevdp • Draft
@jakevdp
Avoid call to asarray in jnp.einsum pull ready Ready for copybara import and testing
#26243 opened Jan 31, 2025 by jakevdp Loading…
@jakevdp
In Progress.
#26242 opened Jan 31, 2025 by copybara-service bot Loading…
[mosaic_gpu/pallas:mgpu] Eradicate wgmma_layout
#26240 opened Jan 31, 2025 by copybara-service bot Loading…
Bump XLA extension version number to fix the build.
#26237 opened Jan 31, 2025 by copybara-service bot Loading…
@emilyfertig
Add FFI example demonstrating the use of XLA's FFI state.
#26235 opened Jan 31, 2025 by copybara-service bot Loading…
@dfm
Remove portpicker for free threaded python 3.13t in test-requirements.txt
#26228 opened Jan 31, 2025 by copybara-service bot Loading…
insane shmap autodiff bug + shit + efficient transpose + eager shmap
#26223 opened Jan 30, 2025 by copybara-service bot Loading…
@yashk2810
Update the calculation for num_processes and num_test_jobs that are used in CUDA test jobs CI Connection Halt - On Retry Flag every job that has its workflow file setup for halting to halt if the workflow is retried
#26192 opened Jan 29, 2025 by copybara-service bot Loading…
@nitins17
Remove srcs_version and python_version attributes, as they already default to "PY3"
#26182 opened Jan 29, 2025 by copybara-service bot Loading…
Remove old backward compatibility mode for old PRGN custom call on GPU
#26179 opened Jan 29, 2025 by copybara-service bot Loading…
@gnecula
linearize the state
#26176 opened Jan 29, 2025 by mattjj Loading…
add jax.distributed.is_initialized
#26172 opened Jan 28, 2025 by ZacCranko Loading…
23
ProTip! Type g i on any issue or pull request to go back to the issue listing page.