Skip to content

Pull requests: google/jax

New pull request New
381 Open 14,141 Closed
381 Open 14,141 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

Update shard_map(jit) to properly set manual_axes on in_shardings and out_shardings of the nested jit. This avoids a problem where the jit returns {manaual} and then this gets passed to ShardToFull (manual is already considered a full sharding).
#22584 opened Jul 23, 2024 by copybara-service bot Loading…
@pschuh
Shmallas, a.k.a. allow lowering shard_map + run_state to a pallas_call.
#22583 opened Jul 23, 2024 by copybara-service bot Loading…
@sharadmv
add test for #17691 pull ready Ready for copybara import and testing
#22582 opened Jul 22, 2024 by mattjj Loading…
1
[xla:cpu] Enable thunks
#22577 opened Jul 22, 2024 by copybara-service bot Loading…
[xla:cpu] Support for up to 16 sorted inputs
#22575 opened Jul 22, 2024 by copybara-service bot Loading…
Remove the unaccelerate_deprecation utility
#22574 opened Jul 22, 2024 by jakevdp Loading…
@jakevdp
3
Reverts d7b821b04d8fec543f570faaece7572a50a75eb6
#22573 opened Jul 22, 2024 by copybara-service bot Loading…
@gnecula
Test sub-host slices, e.g., a TPU v2 1x1 slice.
#22565 opened Jul 22, 2024 by copybara-service bot Loading…
[Mosaic GPU] Replace multicast_mask by a nicer collective async copy interface
#22562 opened Jul 22, 2024 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Add early support for block clusters and multicast TMA
#22561 opened Jul 22, 2024 by copybara-service bot Loading…
@apaszke
Set //third_party/py/jax/tests:lax_numpy_test to use static linking.
#22560 opened Jul 22, 2024 by copybara-service bot Loading…
1
Another take at enabling Pallas GPU tests on x64
#22559 opened Jul 22, 2024 by superbobry • Draft
1
Fixed pl.BlockSpec argument ordering in the Pallas TPU matmul tutorial
#22558 opened Jul 22, 2024 by superbobry Loading…
Show cache key when logging compilation cache hits/misses pull ready Ready for copybara import and testing
#22556 opened Jul 22, 2024 by cool-RR Loading…
1
10
DOC: Improved docs for jnp.fft.fft2 and jnp.fft.ifft2
#22555 opened Jul 22, 2024 by rajasekharporeddy Loading…
[pallas] Simplify GridMapping pull ready Ready for copybara import and testing
#22552 opened Jul 21, 2024 by gnecula Loading…
@gnecula
Reverts 093b92be8ed7bd979486614325956e88cc474ff1
#22551 opened Jul 21, 2024 by copybara-service bot Loading…
@gnecula
[Mosaic TPU] (8,128),-2 -> (8, 128) for non-zero and replicated 2nd minor offset
#22542 opened Jul 20, 2024 by copybara-service bot Loading…
[Mosaic TPU] Enable lane broadcast for packed types and offsets outside of first tile, and fix some broadcast infer logic
#22537 opened Jul 19, 2024 by copybara-service bot Loading…
Add logging the jax2tf mlir_module_serialized module size.
#22536 opened Jul 19, 2024 by copybara-service bot Loading…
Skip some Pallas tests that fail on TPUv6.
#22535 opened Jul 19, 2024 by copybara-service bot Loading…
@hawkinsp
[XLA:Mosaic] Create apply layout pass with ctx instead of config list.
#22528 opened Jul 19, 2024 by copybara-service bot Loading…
[pallas] Add lowering errors for block shapes that are not supported.
#22527 opened Jul 19, 2024 by copybara-service bot Loading…
@gnecula
fix basic scan bug with attrs pull ready Ready for copybara import and testing
#22523 opened Jul 19, 2024 by mattjj Loading…
improve stop_gradient(mutable_array) error pull ready Ready for copybara import and testing
#22521 opened Jul 18, 2024 by mattjj Loading…
3
ProTip! no:milestone will show everything without a milestone.