Skip to content

Pull requests: jax-ml/jax

New pull request New
421 Open 15,998 Closed
421 Open 15,998 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

Document cudaMallocAsync
#25082 opened Nov 25, 2024 by nouiz Loading…
[Mosaic GPU] Simplify logic for pointwise splat operands
#25081 opened Nov 25, 2024 by copybara-service bot Loading…
@apaszke
Use xla_extension_version instead of jaxlib_version
#25079 opened Nov 25, 2024 by copybara-service bot Loading…
[pallas_mgpu] Map indices primitive
#25078 opened Nov 25, 2024 by copybara-service bot Loading…
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…
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
2
Refactor FFI examples to consolidate several examples into one submodule pull ready Ready for copybara import and testing
#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…
ProTip! Adding no:label will show everything without a label.