Skip to content

Pull requests: jax-ml/jax

New pull request New
404 Open 16,632 Closed
404 Open 16,632 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

[Mosaic] Fix infer/apply extensions.
#25868 opened Jan 14, 2025 by copybara-service bot Loading…
Fix remat bug on primitives with multiple outputs.
#25866 opened Jan 14, 2025 by copybara-service bot Loading…
@justinjfu
In progress experimention. Add StringDType to JAX's supported types.
#25865 opened Jan 13, 2025 by copybara-service bot Loading…
More linearization fixes pull ready Ready for copybara import and testing
#25864 opened Jan 13, 2025 by dougalm Loading…
[Docs] Remove --xla_gpu_enable_triton_softmax_fusion from docs
#25863 opened Jan 13, 2025 by copybara-service bot Loading…
@bchetioui
[Mosaic] Allow passing ApplyVectorLayoutCtx to tpu.apply_layout_op.
#25862 opened Jan 13, 2025 by copybara-service bot Loading…
Bump actions/upload-artifact from 4.5.0 to 4.6.0 dependencies Pull requests that update a dependency file github_actions Pull requests that update Github_actions code
#25859 opened Jan 13, 2025 by dependabot bot Loading…
[pallas:mosaic_gpu] Minor fixes regarding the layout types.
#25858 opened Jan 13, 2025 by copybara-service bot Loading…
[Pallas TPU] Remove forward compatibility code for float -> signed conversions
#25854 opened Jan 13, 2025 by copybara-service bot Loading…
@apaszke
[Pallas] Add pipeline mode to pltpu
#25852 opened Jan 12, 2025 by copybara-service bot Loading…
[pallas][mosaic kernel export] Add initial support for exporting a dynamic shapes (placeholder bound) kernel out of mosaic, via pallas as both MLIR and jaxpr.
#25849 opened Jan 12, 2025 by copybara-service bot Loading…
Update ravel_multi_index to accept non-concrete values.
#25840 opened Jan 10, 2025 by copybara-service bot Loading…
Adding GPU paged attention kernel pull ready Ready for copybara import and testing
#25839 opened Jan 10, 2025 by Rifur13 Loading…
jax.vmap: convert mapped input arguments to array
#25835 opened Jan 10, 2025 by jakevdp Loading…
1
@jakevdp
#sdy dynamically choose which custom_partitioning API to use based on the current
#25834 opened Jan 10, 2025 by copybara-service bot Loading…
[Mosaic TPU] Add a faster implementation for packing b16 to s8 in TPUv6
#25829 opened Jan 10, 2025 by copybara-service bot Loading…
@apaszke
[better_errors] Add debug_info to DynamicJaxprTrace and JaxprStackFrame better_errors Improve the error reporting pull ready Ready for copybara import and testing
#25827 opened Jan 10, 2025 by gnecula Loading…
@gnecula
4
[Mosaic TPU] Add support for packing to 16-bit integers on TPUv4
#25824 opened Jan 10, 2025 by copybara-service bot Loading…
@apaszke
[xla:cpu] FFI: Implement Jax callbacks for CPU devices using XLA FFI.
#25815 opened Jan 9, 2025 by copybara-service bot Loading…
[Will not merge] Show the IR of the partial-auto boundary.
#25813 opened Jan 9, 2025 by yliu120 Loading…
1
Add cuDNN flash attention sequence packing
#25812 opened Jan 9, 2025 by Cjkkkk Loading…
1
[mosaic_gpu] Added a serialization pass
#25811 opened Jan 9, 2025 by copybara-service bot Loading…
@superbobry
[ROCm] Add post-merge check on main for AMD GPUs
#25810 opened Jan 9, 2025 by charleshofer Loading…
[better_errors] Ensure that tracer errors in for_loop point to use code pull ready Ready for copybara import and testing
#25798 opened Jan 9, 2025 by gnecula Loading…
1
@gnecula
5
Collective IDs *do* guarantee that we get the same semaphore every time.
#25794 opened Jan 9, 2025 by copybara-service bot Loading…
ProTip! Type g p on any issue or pull request to go back to the pull request listing page.