Skip to content

Pull requests: jax-ml/jax

New pull request New
399 Open 15,067 Closed
399 Open 15,067 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

DOC: Improve docs for jax.numpy: conjugate, conj, imag and real
#23816 opened Sep 21, 2024 by rajasekharporeddy Loading…
Make make_array_from_process_local_data go via device_put if there is only 1 process.
#23815 opened Sep 21, 2024 by copybara-service bot Loading…
@yashk2810
[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…
add jax.experimental.primal_tangent_dtype helper pull ready Ready for copybara import and testing
#23812 opened Sep 20, 2024 by mattjj 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
Update a single name to test tooling post move
#23781 opened Sep 19, 2024 by MichaelHudgins Loading…
[Pallas] Add lowering for threefry PRNG.
#23778 opened Sep 19, 2024 by copybara-service bot Loading…
@justinjfu
ProTip! no:milestone will show everything without a milestone.