Skip to content

Pull requests: jax-ml/jax

New pull request New
386 Open 15,708 Closed
386 Open 15,708 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

Remove a number of expired deprecations. pull ready Ready for copybara import and testing
#24652 opened Oct 31, 2024 by jakevdp Loading…
@jakevdp
[sharding_in_types] Allow device_put to reshard inputs. device_put is a good choice for resharding since it already handles transpose correctly because it tracks the src sharding too.
#24650 opened Oct 31, 2024 by copybara-service bot Loading…
@yashk2810
[Mosaic GPU] Improve correctness of benchmarking scripts
#24645 opened Oct 31, 2024 by andportnoy Loading…
1
#sdy Fix Shardy bug where we weren't setting shmap in/out shardings as open
#24644 opened Oct 31, 2024 by copybara-service bot Loading…
[export] Cleanup in the export module.
#24643 opened Oct 31, 2024 by copybara-service bot Loading…
@gnecula
Verify we can offload more than a single computation to SparseCore
#24638 opened Oct 30, 2024 by copybara-service bot Loading…
Expose equality_errors in tree_util because it's such a useful tool.
#24634 opened Oct 30, 2024 by copybara-service bot Loading…
@IvyZX
[Pallas:MGPU] Add support for multiple heads in attention
#24630 opened Oct 30, 2024 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Fix the ordering of transforms in async_copy
#24629 opened Oct 30, 2024 by copybara-service bot Loading…
@apaszke
Bump hypothesis from 6.102.4 to 6.115.6 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#24610 opened Oct 30, 2024 by dependabot bot Loading…
[NVIDIA] Cudnn dot_product_attention supports head size up to 256
#24607 opened Oct 29, 2024 by kaixih Loading…
GPU Pallas decode_attention improvements
#24602 opened Oct 29, 2024 by rdyro Loading…
Refactor JAX build wheel rule.
#24594 opened Oct 29, 2024 by copybara-service bot Loading…
represent random.key_impl of builtin RNGs by canonical string name pull ready Ready for copybara import and testing
#24593 opened Oct 29, 2024 by froystig Loading…
@froystig
2
[pallas:mosaic_gpu] Ported two pipelining optimizations to emit_pipeline
#24590 opened Oct 29, 2024 by copybara-service bot Loading…
@superbobry
Add a link to the Intel OneAPI plugin for JAX
#24567 opened Oct 28, 2024 by mini-goel Loading…
Plumb dot dimension numbers into TPU matmul op.
#24561 opened Oct 28, 2024 by copybara-service bot Loading…
Extend PR23304 to include jnp.maximum and jnp.minimum
#24551 opened Oct 28, 2024 by copybara-service bot Loading…
[pallas:mosaic_gpu] emit_pipeline now maintains the grid indices
#24546 opened Oct 27, 2024 by copybara-service bot Loading…
@superbobry
Update pre-commit workflow to cache on jax version pull ready Ready for copybara import and testing
#24543 opened Oct 26, 2024 by dfm Loading…
@dfm
Remove the initial argument to jax.nn.softmax and jax.nn.log_softmax.
#24540 opened Oct 25, 2024 by copybara-service bot Loading…
@jakevdp
Limit self-hosted jobs to JAX main repo pull ready Ready for copybara import and testing
#24536 opened Oct 25, 2024 by MichaelHudgins Loading…
Update jaxlib GPU FFI handlers to use templated typed_data method rather than explicit casting.
#24534 opened Oct 25, 2024 by copybara-service bot Loading…
@dfm
improve checkpoint / remat concreteness error with static_argnums pull ready Ready for copybara import and testing
#24516 opened Oct 24, 2024 by mattjj Loading…
Add composite test to jax.
#24515 opened Oct 24, 2024 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.