Skip to content

Pull requests: google/jax

New pull request New
364 Open 13,923 Closed
364 Open 13,923 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

[Pallas TPU] Add lowering for lax.sign
#22283 opened Jul 5, 2024 by ayaka14732 • Draft
@ayaka14732
1
[Pallas TPU] Add lowering for lax.erf_inv 32 bit
#22282 opened Jul 4, 2024 by ayaka14732 Loading…
@ayaka14732
4
[Pallas TPU] Add lowering for lax.shift_right_arithmetic
#22279 opened Jul 4, 2024 by ayaka14732 Loading…
@ayaka14732
[pallas] Fix the interpreter for block_shape not dividing the overall shape pull ready Ready for copybara import and testing
#22275 opened Jul 4, 2024 by gnecula Loading…
@gnecula
1
[pallas] fsync the stdout file descriptor before capturing.
#22273 opened Jul 4, 2024 by copybara-service bot Loading…
Respect the share_binary_between_hosts flag
#22272 opened Jul 4, 2024 by ayaka14732 Loading…
1
@ayaka14732
1
Update persistent compilation cache doc
#22271 opened Jul 4, 2024 by ayaka14732 • Draft
@ayaka14732
Tighten Eigh tolerances back up by after forcing more accurate computation of sin/cosine factors in the two-sided Jacobi eigensolver.
#22267 opened Jul 4, 2024 by copybara-service bot Loading…
Don't wrap singleton ir.Types during HLO lowering. pull ready Ready for copybara import and testing
#22263 opened Jul 3, 2024 by hawkinsp Loading…
Fix C++ registration of FFI handlers and consolidate gpu/linalg kernel implementation.
#22262 opened Jul 3, 2024 by copybara-service bot Loading…
@dfm
DOC: Improved docs for jnp.fft.fftn
#22258 opened Jul 3, 2024 by rajasekharporeddy Loading…
@dfm
7
[JAX] Raise an error if where is not of dtype bool in jax.numpy.mean.
#22256 opened Jul 3, 2024 by copybara-service bot Loading…
Replaced string_stdout with jtu.capture_stdout in tpu/pallas_call_test.py
#22255 opened Jul 3, 2024 by copybara-service bot Loading…
@superbobry
Initial set of changes to allow for lowering to the Shardy dialect.
#22254 opened Jul 3, 2024 by copybara-service bot Loading…
DOC: Improved docs for jnp.ptp and count_nonzero
#22248 opened Jul 3, 2024 by rajasekharporeddy Loading…
@dfm
8
[Pallas] Add better error message for a misspecified BlockSpec
#22247 opened Jul 3, 2024 by copybara-service bot Loading…
@sharadmv
make remat reduce precision of saved values to avoid xla excess precision
#22244 opened Jul 2, 2024 by mattjj • Draft
Extend CustomCallOp backend_config to take a DictionaryAttr
#22242 opened Jul 2, 2024 by copybara-service bot Loading…
jax.numpy: better docstring for array_equiv function
#22241 opened Jul 2, 2024 by selamw1 Loading…
@dfm
4
[Pallas] Add better reduction support.
#22236 opened Jul 2, 2024 by copybara-service bot Loading…
@justinjfu
Enable opt-in autodetection of distributed configuration for mpi4py, attempt 2
#22235 opened Jul 2, 2024 by coreyjadams • Draft
2
[Mosaic_GPU] Allow measured functions to be rerun without recompiling
#22234 opened Jul 2, 2024 by copybara-service bot Loading…
[xla:ffi] Use lazy decoding for Buffer<dtype,rank>
#22224 opened Jul 1, 2024 by copybara-service bot Loading…
Improve documentation of jax.numpy.clip documentation pull ready Ready for copybara import and testing
#22193 opened Jun 28, 2024 by jakevdp Loading…
@jakevdp
1
better log10 doc pull ready Ready for copybara import and testing
#22192 opened Jun 28, 2024 by pkgoogle Loading…
@dfm
5
ProTip! Exclude everything labeled bug with -label:bug.