Skip to content

Pull requests: google/jax

New pull request New
329 Open 13,339 Closed
329 Open 13,339 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

Raise a lowering-time error when broadcasted operand has invalid shape
#21482 opened May 29, 2024 by copybara-service bot Loading…
@superbobry
[shape_poly] Add support for jnp.tril. pull ready Ready for copybara import and testing
#21481 opened May 29, 2024 by gnecula Loading…
@gnecula
[jax2tf] Fix device assignment error for grad under native serialization pull ready Ready for copybara import and testing
#21479 opened May 29, 2024 by gnecula Loading…
1
@gnecula
Make attrs work with pytrees pull ready Ready for copybara import and testing
#21478 opened May 29, 2024 by dougalm Loading…
1
Improve docs for jnp.partition & argpartition documentation pull ready Ready for copybara import and testing
#21476 opened May 29, 2024 by jakevdp Loading…
@jakevdp
Add a link to another page with XLA flags.
#21473 opened May 29, 2024 by nouiz Loading…
For MATT: x64 in jaxpr eqn context
#21472 opened May 28, 2024 by copybara-service bot Loading…
@yashk2810
[array api] update test suite to most recent commit
#21469 opened May 28, 2024 by jakevdp Loading…
@jakevdp
jnp.searchsorted: support sorter argument
#21468 opened May 28, 2024 by jakevdp Loading…
@jakevdp
[Mosaic] Fix bug in VectorLayout::generalizes introduced in cl/636250759
#21464 opened May 28, 2024 by copybara-service bot Loading…
Makes global_shape optional for jax.make_array_from_process_local_data.
#21456 opened May 28, 2024 by copybara-service bot Loading…
[Mosaic GPU] Add a synchronization point at the end of the constructor for BarrierArray.
#21450 opened May 28, 2024 by copybara-service bot Loading…
@bchetioui
Bump zipp from 3.18.1 to 3.19.0 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#21445 opened May 27, 2024 by dependabot bot Loading…
Added trace alias to jnp.linalg
#21442 opened May 27, 2024 by vfdev-5 • Draft
@jakevdp
4
Updated jnp.ceil/floor/trunc to preserve int dtypes
#21441 opened May 27, 2024 by vfdev-5 • Draft
1
Added op name and dtype into testcase name for tests in lax_numpy_operators_test.py
#21440 opened May 27, 2024 by vfdev-5 • Draft
Add reference to Flax NNX
#21439 opened May 27, 2024 by cgarciae Loading…
Support Hessian of gamma-distributed samples
#21432 opened May 25, 2024 by NeilGirdhar Loading…
1
remove is_training from cuDNN SDPA API pull ready Ready for copybara import and testing
#21430 opened May 24, 2024 by Cjkkkk Loading…
4
[Mosaic GPU] Add the option to control output dtype in the matmul examples
#21427 opened May 24, 2024 by copybara-service bot Loading…
[Mosaic GPU] Support for different accumulator types
#21426 opened May 24, 2024 by copybara-service bot Loading…
[Mosaic GPU] Integrate the ExecutionEngine with the jaxlib GPU plugin
#21425 opened May 24, 2024 by copybara-service bot Loading…
@apaszke
[XLA][Mosaic] Add support for fp8 matmuls
#21419 opened May 24, 2024 by copybara-service bot Loading…
DOC: Add code examples to jax.scipy.fft.dct doc
#21418 opened May 24, 2024 by rajasekharporeddy Loading…
@jakevdp
[Mosaic] Packed loads and stores with 1D tiling should use (1, 128 * packing)
#21417 opened May 24, 2024 by copybara-service bot Loading…
ProTip! Updated in the last three days: updated:>2024-05-26.