Pull requests: google/jax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
[pallas] align interpreter load/store with masked behaviour
#21298
opened May 18, 2024 by
oliverdutton
Loading…
Add unit test to check if the backend serialization/deserialization result equal to the original executable.
#21291
opened May 17, 2024 by
copybara-service
bot
Loading…
Finalize deprecation of the config module.
#21288
opened May 17, 2024 by
copybara-service
bot
Loading…
[export] Relax the check that exported modules are used with same number of devices as when exported
pull ready
Ready for copybara import and testing
#21280
opened May 17, 2024 by
gnecula
Loading…
Simplify JAX lowering rules for cumulative sum
#21274
opened May 17, 2024 by
copybara-service
bot
Loading…
Removed noop # type: ignore comments
pull ready
Ready for copybara import and testing
#21273
opened May 17, 2024 by
superbobry
Loading…
[Mosaic] Also check bitwidth in apply-vector-layout's
layoutIsValidForValue
.
#21272
opened May 17, 2024 by
copybara-service
bot
Loading…
[Mosaic GPU] A facility to check for performance regressions.
#21271
opened May 17, 2024 by
copybara-service
bot
Loading…
[Mosaic GPU] Fix the accumulator sync fence by defining the registers instead of using constants.
#21270
opened May 17, 2024 by
copybara-service
bot
Loading…
jax.nn.softmax: only use initial=-inf when necessary
#21264
opened May 16, 2024 by
copybara-service
bot
Loading…
CompatTest.test_cuda_eigh_cusolver_syev_dtype=f32_syevd: skip for CUDA 12.4+
#21258
opened May 16, 2024 by
olupton
Loading…
[Mosaic GPU] Prepare mamul example so it can be exposed to other projects.
#21255
opened May 16, 2024 by
copybara-service
bot
Loading…
[Pallas] Add temporary skip for pallas out-of-bounds interpret mode test on GPU.
#21252
opened May 15, 2024 by
copybara-service
bot
Loading…
Add Ready for copybara import and testing
boundary='wrap'
support for jax.scipy.signal.convolve2d
and jax.scipy.signal.correlate2d
pull ready
#21241
opened May 15, 2024 by
rajasekharporeddy
Loading…
[DOC] Add documentation for jax.experimental.serialize_executable module
pull ready
Ready for copybara import and testing
#21231
opened May 14, 2024 by
nouiz
Loading…
Test the loading of parameters from host memory
#21230
opened May 14, 2024 by
copybara-service
bot
Loading…
Add new flash attn features to cuDNN SDPA API and remove fused attn
pull ready
Ready for copybara import and testing
#21228
opened May 14, 2024 by
Cjkkkk
Loading…
[export] Enable model replication sharding using
jax.pmap
.
#21219
opened May 14, 2024 by
copybara-service
bot
Loading…
Pulled build_mosaic_gpu default into jax.bzl
#21205
opened May 13, 2024 by
copybara-service
bot
Loading…
Previous Next
ProTip!
no:milestone will show everything without a milestone.