-
Couldn't load subscription status.
- Fork 3.2k
Pull requests: jax-ml/jax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
nn: use ensure_arraylike instead of check_arraylike and jnp.asarray
#32990
opened Oct 30, 2025 by
arvindsankar
Loading…
[Mosaic TPU] Add pack/unpack elementwise operations.
#32988
opened Oct 29, 2025 by
copybara-service
bot
Loading…
Added transpose rhs option to ragged dot kernel, useful for implementing the backward function.
#32986
opened Oct 29, 2025 by
copybara-service
bot
Loading…
Use significantly fewer processes for Windows nightlies under pytest.
#32984
opened Oct 29, 2025 by
copybara-service
bot
Loading…
[doc] specify help string by keyword in configs.
documentation
kokoro:force-run
pull ready
Ready for copybara import and testing
#32982
opened Oct 29, 2025 by
jakevdp
Loading…
Use reuse semantics instead of copy semantics when calling DisassembleIntoSingleDeviceArrays in the IFRT proxy, since disassemble always aliases.
#32977
opened Oct 29, 2025 by
copybara-service
bot
Loading…
Rename jax._src.dtypes.dtype to jax._src.dtypes.user_dtype_like_to_dtype.
#32976
opened Oct 29, 2025 by
copybara-service
bot
Loading…
Add support for HiType in Pallas Fuser API
#32975
opened Oct 29, 2025 by
copybara-service
bot
Loading…
[pallas:sc] Lower ref indexing to memref.subview instead of tpu.memref_slice and tpu.memref_squeeze
#32973
opened Oct 29, 2025 by
copybara-service
bot
Loading…
#pallas Avoid scaling the pallas_call CostEstimate by the batching dimension size if it is a dynamic shape.
#32971
opened Oct 29, 2025 by
copybara-service
bot
Loading…
Do not error during type promotion when promoting explicit x64 values in a non-x64 mode.
#32967
opened Oct 29, 2025 by
copybara-service
bot
Loading…
[pallas:sc] Removed unnecessary libtpu version guards in tests
#32959
opened Oct 29, 2025 by
copybara-service
bot
Loading…
[Mosaic GPU] Remove WGMMA tests covering large tiles.
#32956
opened Oct 29, 2025 by
copybara-service
bot
Loading…
Allow type promotion of small integer types.
#32954
opened Oct 29, 2025 by
copybara-service
bot
Loading…
[Mosaic GPU][NFC] Update the documentation of
wgmma.
#32952
opened Oct 29, 2025 by
copybara-service
bot
Loading…
[Pallas TPU] Automatically plumb name argument to pallas_call into name_stack
#32951
opened Oct 29, 2025 by
copybara-service
bot
Loading…
[NFC] Add
name argument to pallas_call primitive.
#32950
opened Oct 28, 2025 by
copybara-service
bot
Loading…
[xla:ffi] Migrate to xla::ffi::MakeTypeInfo() API
#32934
opened Oct 28, 2025 by
copybara-service
bot
Loading…
[hijax] rename NewstyleHiPrimitive, expose in jax.experimental
#32933
opened Oct 28, 2025 by
copybara-service
bot
Loading…
Use a faster pathway for i4 -> f8 canonicalization on TPU 7x: i4 -> i8 -> f8 instead of i4 -> i32 -> f32 -> f8.
#32931
opened Oct 28, 2025 by
copybara-service
bot
Loading…
Previous Next
ProTip!
Add no:assignee to see everything that’s not assigned.