Skip to content

Pull requests: jax-ml/jax

New pull request New
411 Open 15,949 Closed
411 Open 15,949 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

Deprecate several private APIs in jax.lib pull ready Ready for copybara import and testing
#25007 opened Nov 20, 2024 by jakevdp Loading…
@jakevdp
[Mosaic GPU] Improve default kernel name and add option to customize
#25006 opened Nov 20, 2024 by andportnoy Loading…
Mention python 3.13 in docs & package metadata pull ready Ready for copybara import and testing
#25005 opened Nov 20, 2024 by jakevdp Loading…
@jakevdp
Make a direct linearize trace
#25004 opened Nov 20, 2024 by dougalm Loading…
[JAX] Ignore process index when calculating module cache.
#25001 opened Nov 20, 2024 by copybara-service bot Loading…
[pallas:mosaic_gpu] Avoid using multiple indexers in the parallel grid test
#25000 opened Nov 20, 2024 by copybara-service bot Loading…
@superbobry
[sharding_in_types] Make flash_attention forward pass in TPU pallas work nicely with sharding in types. Backward pass is still busted which I will fix in follow up CLs.
#24994 opened Nov 20, 2024 by copybara-service bot Loading…
@yashk2810
Don't psum over auto mesh dims in _unmentioned2.
#24993 opened Nov 20, 2024 by copybara-service bot Loading…
@pschuh
Consolidate material on debugging NaNs.
#24989 opened Nov 19, 2024 by emilyfertig Loading…
3
Add linear layout support to host compute's params/results that live on device.
#24988 opened Nov 19, 2024 by copybara-service bot Loading…
[Mosaic TPU] Add bound check for general vector store op.
#24985 opened Nov 19, 2024 by copybara-service bot Loading…
[Pallas] Use Pallas cost estimator for flash attention.
#24980 opened Nov 19, 2024 by copybara-service bot Loading…
@justinjfu
Delete non-public API jax.lib.xla_bridge._backends
#24976 opened Nov 19, 2024 by copybara-service bot Loading…
@jakevdp
Remove unneeded dependency from rocm_plugin_extension.
#24975 opened Nov 19, 2024 by copybara-service bot Loading…
Run the TPU workflow on new self-hosted runners
#24972 opened Nov 19, 2024 by copybara-service bot Loading…
Removed a few deprecated APIs from jax.lib and jax.core
#24969 opened Nov 19, 2024 by copybara-service bot Loading…
@superbobry
[pallas mgpu] Lowering for while loops as long as they are secretly for loops.
#24965 opened Nov 19, 2024 by copybara-service bot Loading…
Add an option to deactivate automatic cluster detection in jax.distributed.initialize().
#24964 opened Nov 19, 2024 by emilyfertig Loading…
3
In progress. Not ready for review. Approach 2.
#24963 opened Nov 19, 2024 by copybara-service bot Loading…
[Doc] Rename gpu_performance_tips.md to performance_tips.md with new CPU performance tips session
#24961 opened Nov 18, 2024 by louie-tsai Loading…
1
[pallas] Expand shape polymorphism support for Pallas. pull ready Ready for copybara import and testing
#24960 opened Nov 18, 2024 by gnecula Loading…
@gnecula
In progress. Not ready for review.
#24954 opened Nov 18, 2024 by copybara-service bot Loading…
Bump hypothesis from 6.102.4 to 6.119.3 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#24950 opened Nov 18, 2024 by dependabot bot Loading…
Bump fonttools from 4.51.0 to 4.55.0 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#24949 opened Nov 18, 2024 by dependabot bot Loading…
[pallas:mosaic_gpu] Cache IR constants in the lowering
#24944 opened Nov 18, 2024 by copybara-service bot Loading…
@superbobry
ProTip! Follow long discussions with comments:>50.