Skip to content

Pull requests: google/jax

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

[Mosaic] Document tpu.create_subelement_mask
#21582 opened Jun 1, 2024 by copybara-service bot Loading…
Port LU Decomposition to XLA's FFI
#21580 opened Jun 1, 2024 by copybara-service bot Loading…
1
Port Cholesky Factorization to XLA's FFI
#21579 opened Jun 1, 2024 by copybara-service bot Loading…
1
[Mosaic] Remove "support" for MAXNUMF vector reductions
#21578 opened Jun 1, 2024 by copybara-service bot Loading…
[Mosaic] Always use vector native tiling when loading from memory that has a single column of tiles
#21575 opened Jun 1, 2024 by copybara-service bot Loading…
Port Cholesky Factorization to XLA's FFI
#21574 opened Jun 1, 2024 by copybara-service bot Loading…
1
Updating jnp.abs/absolute docs
#21571 opened May 31, 2024 by pkgoogle Loading…
chore: adopt new local wheel installation logic
#21570 opened May 31, 2024 by copybara-service bot Loading…
[jax:cpu] Add a test config to run JAX with new XLA:CPU runtime
#21569 opened May 31, 2024 by copybara-service bot Loading…
skip _origin_msg invar debug info if invar_pos/arg_info is malformed pull ready Ready for copybara import and testing
#21567 opened May 31, 2024 by mattjj Loading…
@mattjj
[JAX] Fail gracefully when an array with multiple shards is passed to make_array_from_single_device_arrays.
#21565 opened May 31, 2024 by copybara-service bot Loading…
@hawkinsp
[Mosaic GPU] Skip tests for Python versions < 3.10
#21560 opened May 31, 2024 by copybara-service bot Loading…
@apaszke
prototyping improving vmap out_axes error
#21559 opened May 31, 2024 by cgarciae Loading…
[jep] Propose a compatibility contract and mechanism for JAX exports
#21550 opened May 31, 2024 by gnecula Loading…
@gnecula
Add loop-based vmap lowering for pallas calls
#21546 opened May 31, 2024 by copybara-service bot Loading…
[Mosaic] In infer-vector-layout's vector.shape_cast rule, remove some obsolete restrictions where batch dims must remain unchanged
#21540 opened May 31, 2024 by copybara-service bot Loading…
[Mosaic] In infer-vector-layout's elementwise rule, when there's a conflict, we can still sometimes keep a replicated offset
#21538 opened May 31, 2024 by copybara-service bot Loading…
relax custom_vjp dtype checks, tweak float0 behavior pull ready Ready for copybara import and testing
#21537 opened May 31, 2024 by mattjj • Draft
[shape_poly] Add limited support for lax.approx_top_k.
#21536 opened May 31, 2024 by copybara-service bot Loading…
@gnecula
[shape_poly] Move to jax._src in preparation for adding to AOT APIs.
#21535 opened May 31, 2024 by copybara-service bot Loading…
@gnecula
[Mosaic] vector.shape_cast sublane (un)folding should support non-zero lane offsets without having to relayout
#21534 opened May 31, 2024 by copybara-service bot Loading…
Use get_api_version to choose gather/scatter
#21533 opened May 31, 2024 by copybara-service bot Loading…
Swap operands of dot if the LHS is fed by a parameter
#21529 opened May 30, 2024 by copybara-service bot Loading…
Rocm dockerfile and test updates
#21528 opened May 30, 2024 by rahulbatra85 Loading…
1
Raise an error when an MLIR lowering is registered for an unknown platform
#21524 opened May 30, 2024 by andportnoy Loading…
14
ProTip! What’s not been updated in a month: updated:<2024-05-02.