Skip to content

Pull requests: jax-ml/jax

New pull request New
419 Open 15,996 Closed
419 Open 15,996 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

Test lax_scipy_spectral_dac error message
#25073 opened Nov 23, 2024 by copybara-service bot Loading…
Move tsl/platform/{build_config,build_config_root,rules_cc}.bzl to xla/tsl/platform
#25072 opened Nov 23, 2024 by copybara-service bot Loading…
[Mosaic TPU] Support packed type matmul with arbitrary shapes.
#25068 opened Nov 22, 2024 by copybara-service bot Loading…
[pallas_mgpu] Allow trees (eg tuples) to be returned from cond_p expressions.
#25063 opened Nov 22, 2024 by copybara-service bot Loading…
[mgpu] Pointwise min
#25061 opened Nov 22, 2024 by copybara-service bot Loading…
[mgpu/pallas_mgpu] Pointwise sign function
#25060 opened Nov 22, 2024 by copybara-service bot Loading…
[mgpu/pallas_mgpu] Pointwise tanh support
#25059 opened Nov 22, 2024 by copybara-service bot Loading…
[pallas:mosaic_gpu] Add test for FragmentedArray.bitcast.
#25058 opened Nov 22, 2024 by copybara-service bot Loading…
@petebu
Added stream annotation support via @compute_on('stream:#') decorator
#25056 opened Nov 22, 2024 by chaserileyroberts Loading…
5
Use optimize='auto' for jnp.linalg.multi_dot
#25055 opened Nov 22, 2024 by dfm Loading…
@dfm
Implement SVD algorithm based on QR for CPU targets
#25053 opened Nov 22, 2024 by JanLuca • Draft
@dfm
1
Rework custom hermetic python instructions.
#25052 opened Nov 22, 2024 by copybara-service bot Loading…
[PjRt-IFRT] Add optional global device mapping support to PjRt-IFRT
#25045 opened Nov 21, 2024 by copybara-service bot Loading…
Add option of disabling cudnn version checking.
#25044 opened Nov 21, 2024 by copybara-service bot Loading…
[ROCM] Jax/Pallas unit test, limiting tile to 64 in for pallas kernel in multihead attention test.
#25043 opened Nov 21, 2024 by zahiqbal Loading…
1
Add an example demonstrating input-output aliasing with the FFI
#25042 opened Nov 21, 2024 by dfm Loading…
1
@dfm
1
Refactor FFI examples to consolidate several examples into one submodule
#25041 opened Nov 21, 2024 by dfm Loading…
@dfm
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0
#25036 opened Nov 21, 2024 by copybara-service bot Loading…
@hawkinsp
Canonicalize platform name for compile-only backends
#25033 opened Nov 21, 2024 by jaro-sevcik Loading…
1
2
Update jnp.floor_divide to make it consistent with np.floor_divide for division by zero
#25032 opened Nov 21, 2024 by rajasekharporeddy Loading…
1
@jakevdp
3
[mgpu] Causal mask operator.
#25028 opened Nov 21, 2024 by copybara-service bot Loading…
[mgpu] A pointwise map function that that also passes the index of the value.
#25027 opened Nov 21, 2024 by copybara-service bot Loading…
[Mosaic] Add extra memref_slice verification and a memory space check helper
#25022 opened Nov 21, 2024 by copybara-service bot Loading…
allow custom_gradient to close over non-autodiff tracers pull ready Ready for copybara import and testing
#25017 opened Nov 20, 2024 by mattjj Loading…
3
[Mosaic GPU] Improve default kernel name and add option to customize pull ready Ready for copybara import and testing
#25006 opened Nov 20, 2024 by andportnoy Loading…
4
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.