Skip to content

Pull requests: google/jax

New pull request New
392 Open 13,022 Closed
392 Open 13,022 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

[pallas] align interpreter load/store with masked behaviour (and adds stride support)
pallas Issues pertaining to Pallas (GPU or TPU)
#21144 opened May 9, 2024 by oliverdutton Loading…
jnp.einsum: improve documentation
documentation
#21138 opened May 8, 2024 by jakevdp Loading…
@jakevdp
Prepare 0.4.28 release.
pull ready Ready for copybara import and testing
#21135 opened May 8, 2024 by hawkinsp Loading…
1
Async dispatch expensive computations on the JAX CPU backend.
#21134 opened May 8, 2024 by copybara-service bot Loading…
@yueshengys
[Mosaic GPU] Cleanup matmul example mainly by removing the m/n block loops (1 step).
#21133 opened May 8, 2024 by copybara-service bot Loading…
1
Simplify JAX lowering rules for cumulative sum
#21132 opened May 8, 2024 by copybara-service bot Loading…
Deprecate newshape argument of jnp.reshape
pull ready Ready for copybara import and testing
#21130 opened May 8, 2024 by Micky774 Loading…
@jakevdp
3
DOC: Add sparse_sigmoid to jax.nn docs
#21125 opened May 8, 2024 by rajasekharporeddy Loading…
Draft of jax.debug.assert
#21124 opened May 7, 2024 by jakevdp • Draft
@jakevdp
jnp.linalg.cond: improve implementation & docs
#21119 opened May 7, 2024 by jakevdp Loading…
@jakevdp
Cholesky rank-1 update kernel for JAX.
#21118 opened May 7, 2024 by copybara-service bot Loading…
1
[XLA:Mosaic] Fix infer/apply vector layout rule for terminators (scf::yieldOp, scf::conditionOp).
#21113 opened May 7, 2024 by copybara-service bot Loading…
Skip float64 test_nextafter on TPU.
pull ready Ready for copybara import and testing
#21108 opened May 7, 2024 by justinjfu Loading…
Fix failing smem_hbm_dma test
kokoro:force-run pull ready Ready for copybara import and testing
#21107 opened May 7, 2024 by justinjfu Loading…
1
jnp.linalg: add precision & preferred_element_type to dot-like functions
pull ready Ready for copybara import and testing
#21106 opened May 7, 2024 by jakevdp Loading…
@jakevdp
[Mosaic GPU] Adjust memref.expand_shape construction to pass in the new args
#21102 opened May 7, 2024 by copybara-service bot Loading…
1
Add jax.make_array_from_process_local_data to create a distributed tensor from host data and supporting scaffolding in sharding to be able to figure out dimensions of host data required.
#21094 opened May 7, 2024 by copybara-service bot Loading…
Remove CreateTokenOp from StableHLO and MHLO, use AfterAllOp
#21093 opened May 7, 2024 by copybara-service bot Loading…
If arguments to jax.jit have shardings on them, add it to avals if jax_enable_sharding_in_avals is on.
#21089 opened May 6, 2024 by copybara-service bot Loading…
@yashk2810
Add support for device kwarg in astype, and add matching utility func
#21086 opened May 6, 2024 by Micky774 Loading…
10
[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
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
[Pallas][Mosaic] Support dynamic roll (part 1) - change roll amount to dynamic
#21063 opened May 3, 2024 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.