Skip to content

Pull requests: google/jax

New pull request New
344 Open 13,532 Closed
344 Open 13,532 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

better docs for right_shift
#21745 opened Jun 7, 2024 by pkgoogle Loading…
@jakevdp
13
Add ffi_call function with a similar signature to pure_callback.
#21743 opened Jun 7, 2024 by copybara-service bot Loading…
@dfm
Enable runtime uptime telemetry for JAX on Cloud TPU.
#21737 opened Jun 7, 2024 by copybara-service bot Loading…
Make pxla.shard_arg batch calls to xc.copy_array_to_devices_with_sharding
#21736 opened Jun 7, 2024 by copybara-service bot Loading…
Updated the type annotations of *_spec= parameters of pl.pallas_call pull ready Ready for copybara import and testing
#21729 opened Jun 7, 2024 by superbobry Loading…
Removed kernel_regeneration_util from Mosaic
#21728 opened Jun 7, 2024 by copybara-service bot Loading…
@superbobry
Activate Cholesky Factorization to XLA's FFI
#21727 opened Jun 7, 2024 by copybara-service bot Loading…
1
DOC: Add example code snippets to jax.scipy.linalg.expm and jax.scipy.linalg.polar docs
#21725 opened Jun 7, 2024 by rajasekharporeddy Loading…
10
[Mosaic GPU] Load LLVM lowering interfaces for all dialects
#21721 opened Jun 7, 2024 by copybara-service bot Loading…
@apaszke
Enable input fusion for a specific kernel pattern.
#21715 opened Jun 6, 2024 by copybara-service bot Loading…
Fix a bug where treedef.flatten_up_to(...) was overly permissive for None treedefs.
#21709 opened Jun 6, 2024 by copybara-service bot Loading…
@hawkinsp
Passes kv_store to ifrt::PjRtClient.
#21707 opened Jun 6, 2024 by copybara-service bot Loading…
@jyingl3
[pallas] Move the hardware_generation query in the code path that needs it
#21706 opened Jun 6, 2024 by copybara-service bot Loading…
@gnecula
[Mosaic GPU] Satisfy MLIR that the dialect interfaces are registered to fix UNDEBUG
#21698 opened Jun 6, 2024 by copybara-service bot Loading…
Migrates multihost_utils to use make_array_from_process_local_data.
#21693 opened Jun 6, 2024 by copybara-service bot Loading…
Refactor & Document
#21691 opened Jun 6, 2024 by Arkay92 Loading…
2
Check that Pallas grid and Mosaic iteration_bounds are positive
#21686 opened Jun 6, 2024 by copybara-service bot Loading…
Allow collectives when only some mesh axes are fully partitioned manually.
#21683 opened Jun 5, 2024 by copybara-service bot Loading…
@pschuh
Rewrite vector.contraction with bf16 accumulator and output into a
#21675 opened Jun 5, 2024 by copybara-service bot Loading…
make cond forward inputs to outputs, reduces vmap lifting pull ready Ready for copybara import and testing
#21673 opened Jun 5, 2024 by mattjj • Draft
Add function to calculate concrete jax2tf output shapes from polymorphic input/output shapes.
#21665 opened Jun 5, 2024 by copybara-service bot Loading…
Added FFI handler registration API to the FFI PjRt extension
#21662 opened Jun 5, 2024 by copybara-service bot Loading…
@superbobry
Moved custom call handler registration into the FFI PjRt extension
#21661 opened Jun 5, 2024 by copybara-service bot Loading…
@superbobry
Add CLI for profiling JAX scripts with output and backend options
#21660 opened Jun 5, 2024 by nirmalmuppiri Loading…
1
Improve tensorstore I/O efficiency
#21654 opened Jun 5, 2024 by copybara-service bot Loading…
ProTip! Exclude everything labeled bug with -label:bug.