Skip to content

Pull requests: jax-ml/jax

New pull request New
392 Open 15,500 Closed
392 Open 15,500 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

Disable breaking test-case
#24380 opened Oct 18, 2024 by copybara-service bot Loading…
Register acos primitive to lower to CHLO acos.
#24378 opened Oct 18, 2024 by copybara-service bot Loading…
Adds arith.select to the op patters in order to canonicalize non 32 bit selects.
#24371 opened Oct 17, 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
nit: Use frozen dataclasses rather than unsafe_hash.
#24369 opened Oct 17, 2024 by copybara-service bot Loading…
@tomhennigan
[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:MGPU] Add support for indexed refs to WGMMA
#24363 opened Oct 17, 2024 by copybara-service bot Loading…
@apaszke
[Pallas:MGPU] Properly commute indexing with other transforms
#24362 opened Oct 17, 2024 by copybara-service bot Loading…
@apaszke
[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
#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…
[pallas] run_scoped now supports partial discharge.
#24344 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
2
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…
Untransform indexing is a Transform protocol method
#24330 opened Oct 16, 2024 by copybara-service bot Loading…
[ROCm] Dockerfile cleanup and dev build fixes
#24321 opened Oct 15, 2024 by mrodden Loading…
1
fix liveness issues within jax itself pull ready Ready for copybara import and testing
#24313 opened Oct 15, 2024 by mattjj Loading…
@mattjj
[Pallas TPU] Add lowerings for bf16 jnp.ceil and jnp.floor in TPU v6+
#24310 opened Oct 15, 2024 by copybara-service bot Loading…
@ayaka14732
Rollback v5e device change
#24304 opened Oct 15, 2024 by copybara-service bot Loading…
resolve extended dtypes in a jaxpr pass
#24302 opened Oct 15, 2024 by froystig • Draft
@froystig @justinjfu
ProTip! Adding no:label will show everything without a label.