Skip to content

Pull requests: google/jax

New pull request New
376 Open 14,272 Closed
376 Open 14,272 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] Enable check for GPU lowering that block sizes are power of 2
#22748 opened Jul 30, 2024 by copybara-service bot Loading…
@gnecula
[pallas] Disallow capturing of consts by kernel functions pull ready Ready for copybara import and testing
#22746 opened Jul 30, 2024 by gnecula Loading…
@gnecula
Fix jax.scipy.stats.betabinom.logpmf and binom.logpmf for JAX to emulate SciPy's behavior when k=n=0
#22744 opened Jul 30, 2024 by rajasekharporeddy Loading…
1
Integrate StableHLO at openxla/stablehlo@fb18ee25
#22743 opened Jul 30, 2024 by copybara-service bot Loading…
Make PjRtCompile accept XlaComputation by value and release XlaComputation as soon as it is converted into HloModule
#22738 opened Jul 29, 2024 by copybara-service bot Loading…
Add optional automatic remat optimization to custom_vjp
#22735 opened Jul 29, 2024 by dfm Loading…
1
Drop support for mhlo in JAX's public API.
#22726 opened Jul 29, 2024 by copybara-service bot Loading…
@hawkinsp
[pallas] Add error message for block_shapes of rank less than 2.
#22722 opened Jul 29, 2024 by copybara-service bot Loading…
@gnecula
Export KeyPath and related types to jax.tree_util
#22720 opened Jul 29, 2024 by copybara-service bot Loading…
Support strided load / store in interpret mode
#22719 opened Jul 29, 2024 by ayaka14732 • Draft
@ayaka14732
1
[pallas] Disallow capturing of consts in the kernel function
#22717 opened Jul 29, 2024 by copybara-service bot Loading…
@gnecula
Add support for bf16 reductions
#22710 opened Jul 29, 2024 by copybara-service bot Loading…
@apaszke
Activate LU Decomposition to XLA's FFI
#22707 opened Jul 29, 2024 by copybara-service bot Loading…
1
[pallas] Improve error localization pull ready Ready for copybara import and testing
#22704 opened Jul 29, 2024 by gnecula Loading…
@gnecula
2
Remove custom partitioning pointer from pre-compiled HLO
#22702 opened Jul 29, 2024 by keshavb96 Loading…
7
Add cudnn_fusion decorator lowering computations to XLA cuDNN fusions.
#22699 opened Jul 27, 2024 by sergachev Loading…
Double buffer pipeline semaphores so we can hide DMA latency under compute and not just BW. Also enable disabling automatic accumulation across pipelines.
#22694 opened Jul 27, 2024 by copybara-service bot Loading…
@epiqueras
[shard-map] relax vmap-of-shmap error with spmd_axis_name pull ready Ready for copybara import and testing
#22676 opened Jul 25, 2024 by mattjj • Draft
Add Scaled Dot Product Attention for FP8
#22670 opened Jul 25, 2024 by wenscarl • Draft
[ROCm]: Add get_arch_details for triton kernel call
#22667 opened Jul 25, 2024 by rahulbatra85 Loading…
1
[shape_poly] Export jax.typing.DimSize and jax.typing.Shape.
#22658 opened Jul 25, 2024 by gnecula Loading…
[ROCM] Fix export harness tests pull ready Ready for copybara import and testing
#22649 opened Jul 25, 2024 by Ruturaj4 Loading…
@gnecula
4
[NVIDIA] Add jit for jax.nn.dot_product_attention tests
#22644 opened Jul 24, 2024 by kaixih Loading…
@gnecula
3
If the product of manual axes is of size 1, then skip emitting
#22641 opened Jul 24, 2024 by copybara-service bot Loading…
@pschuh
Ignore more XLA flags for the JAX cache.
#22624 opened Jul 24, 2024 by nouiz Loading…
1
ProTip! Type g i on any issue or pull request to go back to the issue listing page.