Skip to content

Pull requests: jax-ml/jax

New pull request New
421 Open 17,158 Closed
421 Open 17,158 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

Make device_put resharding on single device array input work under use_mesh. Fixes https://github.com/jax-ml/jax/issues/26552
#26553 opened Feb 15, 2025 by copybara-service bot Loading…
1
@yashk2810
jax.lax: improve docs for bitwise operators. documentation
#26550 opened Feb 14, 2025 by jakevdp Loading…
@jakevdp
Support error checking in auto mode outside jit
#26549 opened Feb 14, 2025 by copybara-service bot Loading…
@ayaka14732
[TPU][MOSAIC][Easy] Add verification for AssumeMultipleOp.
#26546 opened Feb 14, 2025 by copybara-service bot Loading…
[CI] testing change to internal sync
#26545 opened Feb 14, 2025 by copybara-service bot Loading…
@MichaelHudgins
Add version check to jaxlib plugin imports.
#26544 opened Feb 14, 2025 by copybara-service bot Loading…
@dfm
Start adding primitive registration helper functions to lax.linalg.
#26543 opened Feb 14, 2025 by copybara-service bot Loading…
@dfm
Fix the type annotations and don't += a generator (it's confusing)
#26541 opened Feb 14, 2025 by copybara-service bot Loading…
@apaszke
Cache shape property on Mesh.
#26540 opened Feb 14, 2025 by copybara-service bot Loading…
Minor format cleanup
#26539 opened Feb 14, 2025 by copybara-service bot Loading…
@ayaka14732
Adds support for string processing in Colocated Python with Pathways backend.
#26531 opened Feb 14, 2025 by copybara-service bot Loading…
[Mosaic] Several fixes/improvements for the new TPU interpret mode.
#26527 opened Feb 13, 2025 by copybara-service bot Loading…
@jburnim
[sharding_in_types] Set the sharding_in_types config to True. This is a purely internal change and shouldn't affect any public APIs.
#26526 opened Feb 13, 2025 by copybara-service bot Loading…
@yashk2810
Add float4_e2m1fn type support
#26525 opened Feb 13, 2025 by wenscarl Loading…
Add jax.random.multinomial.
#26524 opened Feb 13, 2025 by carlosgmartin Loading…
[Pallas] Reductions with replicated axes.
#26516 opened Feb 13, 2025 by copybara-service bot Loading…
Upgrade to Bazel 7.x
#26513 opened Feb 13, 2025 by copybara-service bot Loading…
Use uv to install Python packages
#26511 opened Feb 13, 2025 by copybara-service bot Loading…
@nitins17
[Mosaic GPU] Add initial support of distributed Pallas on GPUs
#26504 opened Feb 12, 2025 by nvcastet Loading…
Codebeaver/test 723699680 26477
#26492 opened Feb 12, 2025 by CodeBeaverAI Loading…
1
#sdy Add JAX backwards compatibility test.
#26484 opened Feb 12, 2025 by copybara-service bot Loading…
[Mosaic TPU] Replace vrecp with vrecp.approx if the f32 div result is used for packing.
#26479 opened Feb 12, 2025 by copybara-service bot Loading…
Add JaxPr cost analysis scaffolding: cost_analysis.py
#26471 opened Feb 11, 2025 by copybara-service bot Loading…
[Mosaic GPU] Implement lowerings for Tile and Transpose transforms from the MLIR dialect.
#26457 opened Feb 11, 2025 by copybara-service bot Loading…
[Mosaic GPU] Add simple Blackwell Flash Attention kernel
#26445 opened Feb 10, 2025 by andportnoy Loading…
ProTip! What’s not been updated in a month: updated:<2025-01-15.