Skip to content

Pull requests: google/jax

New pull request New
384 Open 12,971 Closed
384 Open 12,971 Closed
Author
Filter by author
Label
Filter by label
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Milestones
Filter by milestone
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

[XLA:TPU] Support output streaming and refactor TryOutputStreaming into a bottoms-up approach.
#21076 opened May 6, 2024 by copybara-service bot Loading…
add error checks for vmap spmd_axis_name
pull ready Ready for copybara import and testing
#21071 opened May 4, 2024 by mattjj Loading…
1
update Jax-Metal CI runner with nightly jaxlib config
#21070 opened May 4, 2024 by shuhand0 Loading…
2
remove named shapes from avals
pull ready Ready for copybara import and testing
#21069 opened May 4, 2024 by mattjj • Draft
[omnitracing] add effectful bind, doesn't skip frames by data dep
#21068 opened May 4, 2024 by mattjj • Draft
[XLA] Add pattern matchers for AsyncStart and AsyncDone.
#21066 opened May 3, 2024 by copybara-service bot Loading…
[Pallas][Mosaic] Support dynamic roll (part 1) - change roll amount to dynamic
#21063 opened May 3, 2024 by copybara-service bot Loading…
[Pallas][Mosaic] Exposes LLO PRNG Ops in Mosaic/Pallas.
#21060 opened May 3, 2024 by copybara-service bot Loading…
@justinjfu
Only check the number of devices match the number at export when > 1.
#21053 opened May 3, 2024 by copybara-service bot Loading…
Remove deprecated kind argument from jnp.sort and jnp.argsort.
#21052 opened May 3, 2024 by copybara-service bot Loading…
@jakevdp
[Mosaic GPU] Prototype FlashAttention kernel with explicit WGMMA/ALU scheduling
#21051 opened May 3, 2024 by copybara-service bot Loading…
@apaszke
add error for remat + rbg for #21046
#21049 opened May 3, 2024 by mattjj • Draft
1
Better documentation for several jax.numpy functions
documentation
#21048 opened May 3, 2024 by jakevdp Loading…
@jakevdp
[Mosaic GPU] Prepare for writing warp-specialized kernels
#21040 opened May 2, 2024 by copybara-service bot Loading…
@apaszke
Add support for TPU delay in Mosaic
#21034 opened May 2, 2024 by copybara-service bot Loading…
1
Add sharding on output avals for shard_map.
#21028 opened May 1, 2024 by copybara-service bot Loading…
@yashk2810
dlpack.callback now supports mutable_results=
#21025 opened May 1, 2024 by copybara-service bot Loading…
@superbobry
Add sharding on ShapedArray and in the type system of jaxpr.
#21022 opened May 1, 2024 by copybara-service bot Loading…
@yashk2810
make _ProfileState's reset point to self.
#21020 opened Apr 30, 2024 by sadikneipp Loading…
[XLA:Mosaic] Support retiling from (8, 128, -2) to (8, 128) for 32-bit data.
#21014 opened Apr 30, 2024 by copybara-service bot Loading…
Adds rewrite patterns for arith and math operations with bf16 operands/results that are not supported by the underlying hardware.
#21007 opened Apr 30, 2024 by copybara-service bot Loading…
1
Added a new callback API which uses DLPack tensors to keep the data on the device
#21003 opened Apr 30, 2024 by copybara-service bot Loading…
@superbobry
Ported threefry2x32 for GPU to the typed XLA FFI
#20997 opened Apr 30, 2024 by copybara-service bot Loading…
@gnecula
[PJRT C API] Enable PJRT C API runtime in jax2tf dlpack.
#20990 opened Apr 29, 2024 by copybara-service bot Loading…
@jyingl3
Generic reduce window jvp
#20987 opened Apr 29, 2024 by copybara-service bot Loading…
13
ProTip! Type g p on any issue or pull request to go back to the pull request listing page.