Skip to content

Pull requests: jax-ml/jax

New pull request New
393 Open 15,390 Closed
393 Open 15,390 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 lax.FftType.
#24230 opened Oct 10, 2024 by copybara-service bot Loading…
@hawkinsp
Remove dead ducc_fft code.
#24229 opened Oct 10, 2024 by copybara-service bot Loading…
@hawkinsp
Add config options to control distributed heartbeat timeout
#24227 opened Oct 10, 2024 by jaro-sevcik Loading…
1
[Pallas TPU] Fix boolean comparison
#24226 opened Oct 10, 2024 by copybara-service bot Loading…
1
@ayaka14732
[sharding_in_types] Normalize partition specs when creating avals so that P(None, None) and P() are treated as replicated and equivalent. Shardings on avals are always normalized.
#24224 opened Oct 10, 2024 by copybara-service bot Loading…
@yashk2810
[Pallas] Add tests for scalar elementwise operations
#24222 opened Oct 10, 2024 by copybara-service bot Loading…
@ayaka14732
[Mosaic TPU] Support native int4 @ int4
#24220 opened Oct 9, 2024 by copybara-service bot Loading…
Update Eigh kernel on GPU to use 64-bit interface when it is available.
#24216 opened Oct 9, 2024 by copybara-service bot Loading…
@dfm
Bump hypothesis from 6.102.4 to 6.114.0 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#24215 opened Oct 9, 2024 by dependabot bot Loading…
Remove deprecated API and formatting fixes.
#24213 opened Oct 9, 2024 by copybara-service bot Loading…
[pallas:mosaic_gpu] pl.run_scoped now supports scoped barriers
#24212 opened Oct 9, 2024 by copybara-service bot Loading…
@superbobry
Activate the FFI implementation of SVD on GPU.
#24211 opened Oct 9, 2024 by copybara-service bot Loading…
@dfm
Deprecate complex dtype support in jnp.arctan2 to make it consistent with np.arctan2
#24203 opened Oct 9, 2024 by rajasekharporeddy Loading…
4
Add config option to log or fatal when jax.Arrays are GCed.
#24202 opened Oct 9, 2024 by copybara-service bot Loading…
[mosaic_gpu] Reshape memref transform.
#24201 opened Oct 9, 2024 by copybara-service bot Loading…
[JAX] Add the function API of jax.experimental.colocated_python
#24200 opened Oct 8, 2024 by copybara-service bot Loading…
[pallas:mosaic_gpu] Added a test doing manual in kernel pipelining
#24198 opened Oct 8, 2024 by copybara-service bot Loading…
@superbobry
Add CI workflow to check JAX distibuted initialize within K8s jobsets
#24197 opened Oct 8, 2024 by yhtang Loading…
Fix listing of vectorized deprecation in changelog.
#24193 opened Oct 8, 2024 by dfm Loading…
@dfm
Migrate CUDA FFI example to GitHub Actions using hosted GPU runner
#24190 opened Oct 8, 2024 by dfm Loading…
@dfm
2
[pallas::mosaic_gpu] re-export mosaic_gpu.tile_shape()
#24188 opened Oct 8, 2024 by copybara-service bot Loading…
[pallas] Added API docs for Triton and Mosaic GPU backends pull ready Ready for copybara import and testing
#24185 opened Oct 8, 2024 by superbobry Loading…
[jax2tf] Disable jax2tf with non-native serialization.
#24183 opened Oct 8, 2024 by copybara-service bot Loading…
@gnecula
Update tf.lite to use ai_edge_litert
#24179 opened Oct 7, 2024 by ecalubaquib Loading…
1
Update JAX release scripts with the new .bazelrc changes
#24178 opened Oct 7, 2024 by copybara-service bot Loading…
ProTip! Follow long discussions with comments:>50.