Skip to content

Pull requests: google/jax

New pull request New
287 Open 8,377 Closed
287 Open 8,377 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

[jax2tf] Add xmap tests
#14182 opened Jan 27, 2023 by copybara-service bot Loading…
Fix the warning being raised when jax.Array is True about using jax.Array
#14177 opened Jan 27, 2023 by copybara-service bot Loading…
register a backend factory for every pjrt plugin found
#14176 opened Jan 27, 2023 by copybara-service bot Loading…
DOC: fix autodidax
pull ready Ready for copybara import and testing
#14175 opened Jan 27, 2023 by jakevdp Loading…
1
@jakevdp
Add PJRT C API to Cloud TPU test matrix
#14174 opened Jan 27, 2023 by skye Loading…
[jax2tf] port sharding tests to use MHLO
#14172 opened Jan 27, 2023 by copybara-service bot Loading…
DOC: clarify purpose of the jax-array doc
#14170 opened Jan 26, 2023 by jakevdp Loading…
@jakevdp
10
Removed deprecated polar_unitary as per phawkins comment.
#14169 opened Jan 26, 2023 by tttc3 Loading…
1
Remove tf-lite deps from tensorflow_core.
#14168 opened Jan 26, 2023 by copybara-service bot Loading…
Make the implementation of Log/Log2/Ln/Log1p in XLA accurate on all platforms, and tighten test tolerances accordingly. Add tests for +/-0, +/-inf, and NaN.
#14142 opened Jan 24, 2023 by copybara-service bot Loading…
1
[XLA:CPU] simple_orc_jit: register PerfJIT listener
#14118 opened Jan 23, 2023 by copybara-service bot Loading…
[shape_poly] Add better support for division, and working with strides
pull ready Ready for copybara import and testing
#14108 opened Jan 21, 2023 by gnecula Loading…
@gnecula
Publish docs for custom_partitioning.
#14071 opened Jan 19, 2023 by copybara-service bot Loading…
@pschuh
Add device_index to hcb.call to configure on which device to invoke host callback in pjit.
#14051 opened Jan 18, 2023 by copybara-service bot Loading…
@zhangqiaorjc
TensorSpec experimental_get_compiler_ir improve the captured_input support.
#14048 opened Jan 18, 2023 by copybara-service bot Loading…
[jax2tf] Fixes a bug in flax model testing.
#14023 opened Jan 16, 2023 by copybara-service bot Loading…
Eagerly register backend factories for PJRT C API plugins.
#14011 opened Jan 14, 2023 by stellaraccident Loading…
@skye
8
Fixed issue with complex-valued inputs to LOBPCG
#14005 opened Jan 13, 2023 by jemisjoky Loading…
1
@jakevdp
3
Added symbolic zeros as inputs to custom_jvp
#13993 opened Jan 12, 2023 by patrick-kidger Loading…
@mattjj
Allow pjit to be picklable.
#13971 opened Jan 11, 2023 by copybara-service bot Loading…
@pschuh
Revive json exporter using a direct XSpace -> json converter (adapted from the
#13968 opened Jan 11, 2023 by copybara-service bot Loading…
@pschuh
checkify: in traceable, convert constvars to args to fit ClosedJaxpr convention
kokoro:force-run pull ready Ready for copybara import and testing
#13966 opened Jan 11, 2023 by mattjj Loading…
@mattjj
Make sure jax.grad(pjit(f)) hits the compilation cache.
#13965 opened Jan 11, 2023 by copybara-service bot Loading…
@apaszke
lax.log: return NaN for out-of-bound gradient
#13945 opened Jan 10, 2023 by jakevdp • Draft
@jakevdp
compute jit number of replicas using least-common-multiple, not max
kokoro:force-run pull ready Ready for copybara import and testing
#13938 opened Jan 10, 2023 by mattjj Loading…
1
@mattjj
ProTip! Adding no:label will show everything without a label.