Skip to content

Pull requests: jax-ml/jax

New pull request New
400 Open 15,052 Closed
400 Open 15,052 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

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…
[Pallas GPU] Enable Pallas OpsExtraTest in 64-bit mode
#23798 opened Sep 20, 2024 by copybara-service bot Loading…
@ayaka14732
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
Better doc for jax.numpy.i0
#23790 opened Sep 20, 2024 by rajasekharporeddy Loading…
13
[host_callback] Adjust the legacy compatibility mode for call_with_device
#23789 opened Sep 20, 2024 by copybara-service bot Loading…
@gnecula
DOC: Improve docs for jax.numpy: deg2rad, rad2deg, degrees, radians
#23788 opened Sep 20, 2024 by rajasekharporeddy Loading…
@jakevdp
6
Add use_cuda_nvcc flag to build.py to enable compilation of CUDA code with clang. If --use_cuda_nvcc NVCC compiler will be used to build CUDA code (default case), if --nouse_cuda_nvcc then Clang will be used to build CUDA code.
#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
Add use_cuda_nvcc flag to build.py to enable compilation of CUDA code with clang. If --use_cuda_nvcc NVCC compiler will be used to build CUDA code (default case), if --nouse_cuda_nvcc then Clang will be used to build CUDA code.
#23777 opened Sep 19, 2024 by copybara-service bot Loading…
Add docstrings for jnp.blackman, jnp.bartlett, jnp.hamming, jnp.hanning, jnp.kaiser
#23776 opened Sep 19, 2024 by copybara-service bot Loading…
@jakevdp
improve tree mismatch error in tree_map
#23771 opened Sep 19, 2024 by mattjj • Draft
2 tasks
1
Remove deprecated partial eval functions (prep work for stackless)
#23768 opened Sep 19, 2024 by copybara-service bot Loading…
@dougalm
Added is_signed to mgpu.FragmentedArray
#23764 opened Sep 19, 2024 by copybara-service bot Loading…
@superbobry
[Mosaic GPU] Clean up the module structure
#23760 opened Sep 19, 2024 by copybara-service bot Loading…
@apaszke
[Pallas] Fix integer array indexing
#23758 opened Sep 19, 2024 by copybara-service bot Loading…
1
@ayaka14732
[JAX] Add test-case on AOT compilation on the same process
#23752 opened Sep 19, 2024 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.