-
Notifications
You must be signed in to change notification settings - Fork 2.6k
Pull requests: google/jax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
Make Tracers unhashable since it's a jax.Array and jax Arrays are unhashable. Fixes https://github.com/google/jax/issues/21824
#21826
opened Jun 12, 2024 by
copybara-service
bot
Loading…
Add frontend attributes to Jax. This allows Jax users to annotate Jax code with frontend_attributes which can be traced down to the HLO level, to be used for numerical debugging purposes. Frontend attributes annotations for instructions added by grad are not supported yet.
#21818
opened Jun 12, 2024 by
copybara-service
bot
Loading…
[PJRT:PLUGIN] Support both string and bytes as the input type of function name for register_custom_call_target in jax-cuda-plugin.
#21814
opened Jun 11, 2024 by
copybara-service
bot
Loading…
Enable jax's cloud-tpu configs when libtpu is present via through "pip install" or set by custom through the $TPU_LIBRARY_PATH env var
#21812
opened Jun 11, 2024 by
copybara-service
bot
Loading…
Add freshness metablock to JAX OSS docs.
#21810
opened Jun 11, 2024 by
copybara-service
bot
Loading…
[ROCm]: Pallas updates for ROCm
pull ready
Ready for copybara import and testing
#21809
opened Jun 11, 2024 by
rahulbatra85
Loading…
[Mosaic GPU] Change FlashAttention implementation to support Grouped Query Attention.
#21808
opened Jun 11, 2024 by
copybara-service
bot
Loading…
prune cond passthrough outputs
pull ready
Ready for copybara import and testing
#21806
opened Jun 11, 2024 by
cgarciae
Loading…
DOC: Better docs for
jnp.convolve
& jnp.correlate
#21804
opened Jun 11, 2024 by
rajasekharporeddy
Loading…
[pallas] Add support for cross-platform lowering
pull ready
Ready for copybara import and testing
#21799
opened Jun 11, 2024 by
gnecula
Loading…
doc: don't use typehints for signatures.
documentation
#21797
opened Jun 11, 2024 by
jakevdp
Loading…
[Mosaic] Simplify vector.shape_cast rules and cover more cases
#21796
opened Jun 11, 2024 by
copybara-service
bot
Loading…
Pallas bitwise_left_shift unit test fix
pull ready
Ready for copybara import and testing
#21780
opened Jun 10, 2024 by
rahulbatra85
Loading…
Rewrite
vector.multi_dim_reduction
with bf16 source/accumulator/output into
#21775
opened Jun 10, 2024 by
copybara-service
bot
Loading…
[export] Add documentation for jax.export
pull ready
Ready for copybara import and testing
#21769
opened Jun 10, 2024 by
gnecula
Loading…
Batch
pxla.shard_args
calls triggered by jax.device_put
#21754
opened Jun 9, 2024 by
copybara-service
bot
Loading…
Add
ffi_call
function with a similar signature to pure_callback
.
#21743
opened Jun 7, 2024 by
copybara-service
bot
Loading…
Enable runtime uptime telemetry for JAX on Cloud TPU.
#21737
opened Jun 7, 2024 by
copybara-service
bot
Loading…
Make
pxla.shard_arg
batch calls to xc.copy_array_to_devices_with_sharding
#21736
opened Jun 7, 2024 by
copybara-service
bot
Loading…
Previous Next
ProTip!
no:milestone will show everything without a milestone.