Skip to content

Pull requests: jax-ml/jax

New pull request New
414 Open 16,720 Closed
414 Open 16,720 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 test case for jit compile of vmap on gpu
#25980 opened Jan 18, 2025 by codinglover222 Loading…
Add reshard API in experimental. Currently for sharding_in_types we have 2 APIs: mesh_cast and reshard. Both work in sharding_in_types mode and affect the sharding of the aval. Following are the semantics of both:
#25978 opened Jan 18, 2025 by copybara-service bot Loading…
@yashk2810
fix aliased arg in async pallas docs
#25977 opened Jan 18, 2025 by arvoelke Loading…
add f-string to mosaic memory space error msg pull ready Ready for copybara import and testing
#25976 opened Jan 18, 2025 by arvoelke Loading…
1
Remove unnecessary use of xla_client.OpMetadata class.
#25975 opened Jan 18, 2025 by copybara-service bot Loading…
@hawkinsp
internal: check integer overflow in lax.asarray pull ready Ready for copybara import and testing
#25969 opened Jan 17, 2025 by jakevdp Loading…
@jakevdp
Use ensure_arraylike utility in jax.numpy sources
#25966 opened Jan 17, 2025 by copybara-service bot Loading…
@jakevdp
[Mosaic GPU] Add manual consumed barrier handling to WS pipeline.
#25964 opened Jan 17, 2025 by copybara-service bot Loading…
@justinjfu
Add DCE rule for custom_jvp
#25963 opened Jan 17, 2025 by dfm Loading…
@dfm
[CI] Add testing for tpu-v6e CI Connection Halt - On Retry Flag every job that has its workflow file setup for halting to halt if the workflow is retried
#25959 opened Jan 17, 2025 by MichaelHudgins • Draft
Add an experimental interface for customizing DCE behavior.
#25956 opened Jan 17, 2025 by copybara-service bot Loading…
@dfm
Enable pivoted QR on GPU via MAGMA.
#25955 opened Jan 17, 2025 by tttc3 Loading…
1
@dfm
[Mosaic TPU] Relax test restrictions after improving DMA support for different tilings
#25954 opened Jan 17, 2025 by copybara-service bot Loading…
@apaszke
WIP: find tests for debug info
#25947 opened Jan 17, 2025 by copybara-service bot Loading…
@gnecula
[Mosaic][TPU] Public desc omitted.
#25946 opened Jan 17, 2025 by copybara-service bot Loading…
Remove all MHLO uses, replace it with StableHLO
#25942 opened Jan 17, 2025 by copybara-service bot Loading…
[Mosaic] Remove hardcoded TARGET_SHAPE and align Python/C++ APIs.
#25939 opened Jan 17, 2025 by copybara-service bot Loading…
[Mosaic TPU] Support non-32 bit mask relayout.
#25938 opened Jan 17, 2025 by copybara-service bot Loading…
[Mosaic TPU] Emulate converting x16 vector to mask if mask packing is supported.
#25937 opened Jan 16, 2025 by copybara-service bot Loading…
Add part (non-quantized K/V pages) of paged_attention_kernel tests back for TPU v6.
#25934 opened Jan 16, 2025 by copybara-service bot Loading…
[jax:custom_partitioning] Extend the custom partitioning API to accept a
#25927 opened Jan 16, 2025 by copybara-service bot Loading…
Add support for constants in the decomposition of lax.composites.
#25921 opened Jan 16, 2025 by copybara-service bot Loading…
@dfm
Adds a minimal but viable implementation of string arrays (with numpy.dtypes.StringDType) in JAX. Currently this only supports making of a string array by means of either jax.numpy.asarray or jax.device_put and reading it back with jax.device_get.
#25918 opened Jan 16, 2025 by copybara-service bot Loading…
[better_errors] Improvements in propagation of debugging info pull ready Ready for copybara import and testing
#25916 opened Jan 15, 2025 by gnecula Loading…
@gnecula
wip sharding in types with yash
#25915 opened Jan 15, 2025 by mattjj • Draft
ProTip! no:milestone will show everything without a milestone.