Skip to content

Pull requests: jax-ml/jax

New pull request New
459 Open 16,496 Closed
459 Open 16,496 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

Add mode='fan_geo_avg' to nn.initializers.variance_scaling.
#25766 opened Jan 8, 2025 by carlosgmartin Loading…
1
[JAX] Add a new jax_num_cpu_devices flag that allows the user to specify the number of CPU directly.
#25765 opened Jan 8, 2025 by copybara-service bot Loading…
@hawkinsp
[Mosaic GPU] Allow multiple gmem indexers on copies.
#25763 opened Jan 7, 2025 by copybara-service bot Loading…
@justinjfu
[Mosaic GPU] Enable loop carries in the pipeline emitter.
#25762 opened Jan 7, 2025 by copybara-service bot Loading…
@justinjfu
[Mosaic] Fix inferMemRefLayout to error out if shape not aligned to tiling.
#25758 opened Jan 7, 2025 by copybara-service bot Loading…
Introduce jax.shard_map, without requiring mesh arg
#25757 opened Jan 7, 2025 by copybara-service bot Loading…
@yashk2810
[ROCm] Implement RNN support
#25755 opened Jan 7, 2025 by Ruturaj4 Loading…
1
[Mosaic GPU] Use num_q_heads=2 in flash_attention.py
#25754 opened Jan 7, 2025 by andportnoy Loading…
Add sph_harm_y to jax.scipy.special and deprecate sph_harm
#25753 opened Jan 7, 2025 by dfm Loading…
@dfm
1
Removed leftover mentions of xmap from the code
#25752 opened Jan 7, 2025 by copybara-service bot Loading…
@superbobry
[shape_poly] Remove the deprecated PolyShape object for specifying symbolic dimensions
#25751 opened Jan 7, 2025 by copybara-service bot Loading…
@gnecula
Add JAX events that have time spans, not only durations.
#25747 opened Jan 7, 2025 by copybara-service bot Loading…
jnp.linalg.solve: finalize deprecation of batched 1D solves pull ready Ready for copybara import and testing
#25741 opened Jan 6, 2025 by jakevdp Loading…
@jakevdp
[Mosaic] Use tpu::CreateMask for getX32VmaskByPaddingEnd.
#25737 opened Jan 6, 2025 by copybara-service bot Loading…
[Mosaic TPU][NFC] Remove redundant num_subelems attribute from CreateSubelementMaskOp
#25729 opened Jan 6, 2025 by copybara-service bot Loading…
[Mosaic TPU] Enable non-sublane-aligned 2D int8 load/stores
#25727 opened Jan 6, 2025 by copybara-service bot Loading…
[Mosaic TPU] Enable unaligned bf16 2D load/stores for earlier TPU gens
#25726 opened Jan 6, 2025 by copybara-service bot Loading…
[Pallas] Add empty/empty_like helper functions
#25722 opened Jan 3, 2025 by copybara-service bot Loading…
@sharadmv
[Mosaic] Create a stub for TPUExtDialect.
#25719 opened Jan 3, 2025 by copybara-service bot Loading…
Add GitHub action workflow for Bazel CUDA continuous tests
#25717 opened Jan 3, 2025 by copybara-service bot Loading…
@nitins17
[Shmap/PartialAuto] Temporary solution for debug.print inside a partial-auto shard map.
#25705 opened Jan 1, 2025 by yliu120 Loading…
Add jax.random.multinomial. kokoro:force-run pull ready Ready for copybara import and testing
#25688 opened Dec 27, 2024 by carlosgmartin Loading…
1
@jakevdp
11
Expand test case to repro a crash.
#25670 opened Dec 23, 2024 by jburnim Loading…
Add support for axis_name and axis_index_groups to lax.ragged_all_to_all
#25660 opened Dec 23, 2024 by copybara-service bot Loading…
Simplify implementation of random.orthogonal.
#25655 opened Dec 22, 2024 by carlosgmartin Loading…
ProTip! Updated in the last three days: updated:>2025-01-04.