Skip to content

Pull requests: jax-ml/jax

New pull request New
401 Open 15,069 Closed
401 Open 15,069 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

[Pallas] Skip atomic_cas and atomic_counter tests on GPU in 64-bit mode
#23828 opened Sep 22, 2024 by copybara-service bot Loading…
@ayaka14732
Simplify extended dtype convert logic pull ready Ready for copybara import and testing
#23823 opened Sep 21, 2024 by mattjj Loading…
1 task
@mattjj
11
Test that jax.numpy docstrings include examples documentation
#23821 opened Sep 21, 2024 by jakevdp Loading…
@jakevdp
Automated Code Change
#23817 opened Sep 21, 2024 by copybara-service bot Loading…
DOC: Improve docs for jax.numpy: conjugate, conj, imag and real documentation pull ready Ready for copybara import and testing
#23816 opened Sep 21, 2024 by rajasekharporeddy Loading…
@jakevdp
3
[Pallas] Simplify sign and erf_inv tests
#23814 opened Sep 21, 2024 by copybara-service bot Loading…
@ayaka14732
[Mosaic TPU] Support bitcast without forcing retiling.
#23813 opened Sep 20, 2024 by copybara-service bot Loading…
Ignore scipy.stats._axis_nan_policy.SmallSampleWarning for LaxBackedScipyStatsTests.testMode
#23810 opened Sep 20, 2024 by copybara-service bot Loading…
Add random.roberts_sequence.
#23808 opened Sep 20, 2024 by carlosgmartin Loading…
[NVIDIA] Fix unsupported dbias in bwd batcher of cudnn sdpa API
#23807 opened Sep 20, 2024 by kaixih Loading…
[ffi] Support handler bundles in GPU plugin extension
#23806 opened Sep 20, 2024 by gspschmid Loading…
@dfm
1
Add an example demonstrating how FFI calls can use global state
#23805 opened Sep 20, 2024 by dfm Loading…
Don't pass --nocheck_visibility to Bazel.
#23804 opened Sep 20, 2024 by hawkinsp Loading…
Better docs for jnp.meshgrid documentation
#23801 opened Sep 20, 2024 by jakevdp Loading…
@jakevdp
[kfac-jax] Update _Linear layer to support the new "algorithm" tuning parameters for dot_general that will be included in the next JAX release.
#23800 opened Sep 20, 2024 by copybara-service bot Loading…
@dfm
#sdy add JAX Shardy support for shard_map.
#23799 opened Sep 20, 2024 by copybara-service bot Loading…
Make an AxisData struct that bundles axis name, size, and spmd name.
#23796 opened Sep 20, 2024 by copybara-service bot Loading…
@dougalm
[pallas::mosaic_gpu] Turn the accumulator into a reference
#23793 opened Sep 20, 2024 by copybara-service bot Loading…
[Mosaic TPU] Improve vector layout inference for vector.shape_cast
#23792 opened Sep 20, 2024 by copybara-service bot Loading…
@apaszke
[Mosaic TPU] Add better tests for newly supported shape_casts (that require a change in tiling)
#23791 opened Sep 20, 2024 by copybara-service bot Loading…
@apaszke
[host_callback] Adjust the legacy compatibility mode for call_with_device
#23789 opened Sep 20, 2024 by copybara-service bot Loading…
@gnecula
Add --use_cuda_nvcc flag to build.py for enable/disable compilation of CUDA by NVCC.
#23787 opened Sep 20, 2024 by copybara-service bot Loading…
[Pallas] Add state discharge rule for pallas_call
#23786 opened Sep 20, 2024 by copybara-service bot Loading…
@sharadmv
Fix errata in block-sparse kernel tutorial.
#23785 opened Sep 20, 2024 by copybara-service bot Loading…
nested jit lower to memory kind properly
#23784 opened Sep 20, 2024 by copybara-service bot Loading…
@yashk2810
ProTip! Updated in the last three days: updated:>2024-09-19.