Skip to content

Pull requests: jax-ml/jax

New pull request New
398 Open 15,514 Closed
398 Open 15,514 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

Add an example for executing serialized function in cpp
#24413 opened Oct 20, 2024 by fenghuizhang Loading…
8
Better docs for jnp.fromfunction documentation
#24412 opened Oct 20, 2024 by jakevdp Loading…
@jakevdp
Clean up conditional lowering of tan after JAX v0.4.34 release. pull ready Ready for copybara import and testing
#24405 opened Oct 19, 2024 by dfm Loading…
@dfm
Better documentation for jnp.load documentation
#24403 opened Oct 19, 2024 by jakevdp Loading…
@jakevdp
2
Make jnp.subtract a ufunc documentation
#24400 opened Oct 19, 2024 by jakevdp Loading…
@jakevdp
Avoid querying metadata query to check if it's GCE if TPU_SKIP_MDS_QUERY is set.
#24394 opened Oct 18, 2024 by copybara-service bot Loading…
Fix env check order by splitting TPU envs into separate files.
#24393 opened Oct 18, 2024 by copybara-service bot Loading…
Respect dot algorithm spec on TPU backends.
#24392 opened Oct 18, 2024 by copybara-service bot Loading…
Document jax config to disable remat HLO pass
#24391 opened Oct 18, 2024 by keshavb96 • Draft
1
[Pallas:MGPU] Add support for passing in WGMMA lhs from registers
#24389 opened Oct 18, 2024 by copybara-service bot Loading…
@apaszke
Add an extension mechanism to run_state that allows:
#24386 opened Oct 18, 2024 by copybara-service bot Loading…
@sharadmv
[Pallas:MGPU] Allow indexing to appear anywhere in the list of transforms
#24383 opened Oct 18, 2024 by copybara-service bot Loading…
@apaszke
Register acos primitive to lower to CHLO acos.
#24378 opened Oct 18, 2024 by copybara-service bot Loading…
Expand the ffi_call API to support more use cases
#24370 opened Oct 17, 2024 by dfm Loading…
@dfm
3
[pallas] input_output_aliases now only include refs which have been written to
#24368 opened Oct 17, 2024 by copybara-service bot Loading…
@superbobry
[Pallas TPU] Fix lowering for jnp.remainder
#24356 opened Oct 17, 2024 by copybara-service bot Loading…
1
@ayaka14732
[Mosaic TPU] Support dynamic DMA and ref slice on the 2nd minor when memref is untiled
#24354 opened Oct 17, 2024 by copybara-service bot Loading…
[Mosaic TPU] Use native vector tiling to load and store with untiled memref.
#24353 opened Oct 17, 2024 by copybara-service bot Loading…
[NVIDIA] Support custom dtype convert in jax.nn.dot_product_attention
#24352 opened Oct 16, 2024 by kaixih Loading…
Improve error message for grad-of-eig better_errors Improve the error reporting
#24350 opened Oct 16, 2024 by jakevdp Loading…
@jakevdp
Move converter related tflite functions to tensorflow/lite repo
#24346 opened Oct 16, 2024 by copybara-service bot Loading…
[export] Add support for serialization for some custom PyTree nodes pull ready Ready for copybara import and testing
#24342 opened Oct 16, 2024 by gnecula Loading…
@gnecula
17
Update ad.py
#24337 opened Oct 16, 2024 by sanowl Loading…
@jakevdp
2
[Pallas TPU] [WIP] Boolean comparison lowering test
#24333 opened Oct 16, 2024 by copybara-service bot Loading…
@ayaka14732
[pallas:mosaic_gpu] Tile and untile of references.
#24331 opened Oct 16, 2024 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.