Skip to content

Pull requests: jax-ml/jax

New pull request New
394 Open 15,751 Closed
394 Open 15,751 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

[WIP] A minimal public API for custom primitives
#24726 opened Nov 5, 2024 by dfm • Draft
@dfm
[Mosaic] Verify that tpu.sem_wait semaphore rank is zero
#24725 opened Nov 5, 2024 by copybara-service bot Loading…
[SDY] add JAX lowering to Shardy ShardingGroupOp for shard_alike.
#24724 opened Nov 5, 2024 by copybara-service bot Loading…
scipy.special.beta: remove deprecated x and y parameters pull ready Ready for copybara import and testing
#24723 opened Nov 5, 2024 by jakevdp Loading…
@jakevdp
Make Tracers non-hashable pull ready Ready for copybara import and testing
#24722 opened Nov 5, 2024 by jakevdp Loading…
1
@jakevdp
Implement the extension to the custom_partitioning API.
#24719 opened Nov 5, 2024 by copybara-service bot Loading…
[Mosaic GPU] Implement tiled and swizzled transfers for tiled layouts
#24718 opened Nov 5, 2024 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Introduce a more flexible layout system
#24714 opened Nov 5, 2024 by copybara-service bot Loading…
@apaszke
[XLA] Make our LLVM usage more googley
#24706 opened Nov 4, 2024 by copybara-service bot Loading…
RFC: specify jit static args via Static annotation enhancement New feature or request
#24705 opened Nov 4, 2024 by jakevdp • Draft
@jakevdp
3
Update link to algebraic_simplifier.cc to point to OpenXLA instead of TF
#24704 opened Nov 4, 2024 by andportnoy Loading…
[jax:custom_partitioning] Implement SdyShardingRule to support
#24701 opened Nov 4, 2024 by copybara-service bot Loading…
register_dataclass: allow marking static fields via field(static=True)
#24700 opened Nov 4, 2024 by jakevdp Loading…
@jakevdp
Add a TPU kernel for eigh.
#24698 opened Nov 4, 2024 by copybara-service bot Loading…
Bump hypothesis from 6.102.4 to 6.116.0 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#24696 opened Nov 4, 2024 by dependabot bot Loading…
Re-factor build CLI to a subcommand based approach
#24695 opened Nov 4, 2024 by copybara-service bot Loading…
[mosaic_gpu] Scalar arguments to kernels.
#24674 opened Nov 2, 2024 by copybara-service bot Loading…
[Mosaic GPU] Implement the skeleton of a lowering pass for the Mosaic GPU dialect.
#24670 opened Nov 1, 2024 by copybara-service bot Loading…
@bchetioui
RFC: add pytree-compatible dataclass decorator
#24664 opened Nov 1, 2024 by jakevdp • Draft
@jakevdp
27
Add a GPU implementation of lax.linalg.eig.
#24663 opened Nov 1, 2024 by copybara-service bot Loading…
@dfm
1
Fix argmin docstring to not say "maximum" pull ready Ready for copybara import and testing
#24657 opened Nov 1, 2024 by dymil Loading…
@dfm
3
Automated Code Change
#24655 opened Nov 1, 2024 by copybara-service bot Loading…
[Mosaic GPU] Improve correctness of benchmarking scripts
#24645 opened Oct 31, 2024 by andportnoy Loading…
1
[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
ProTip! no:milestone will show everything without a milestone.