Pull requests: jax-ml/jax
Pull requests list
docstrings and API reference doc listing for the traced AOT stage
kokoro:force-run
pull ready
Ready for copybara import and testing
#26478
opened Feb 12, 2025 by
froystig
Loading…
[sharding_in_types] When caching mesh with axis_types, make sure the data structure is (axis_size, axis_names, tuple(axis_types))
#26473
opened Feb 12, 2025 by
copybara-service
bot
Loading…
refactor: move jnp.einsum impl into its own submodule
#26472
opened Feb 11, 2025 by
jakevdp
Loading…
[Test] [Do not review yet] Add core logic of the
error_check module
#26469
opened Feb 11, 2025 by
copybara-service
bot
Loading…
Only cache jax.Array._npy_value when a copy is required.
#26460
opened Feb 11, 2025 by
copybara-service
bot
Loading…
[better_errors] Continue adding debug info to Jaxprs (step 8)
pull ready
Ready for copybara import and testing
#26455
opened Feb 11, 2025 by
gnecula
Loading…
[better_errors] Make it explicit that debug_info is not None.
#26451
opened Feb 11, 2025 by
copybara-service
bot
Loading…
refactor: move lax_numpy tensor contractions into their own file
pull ready
Ready for copybara import and testing
#26447
opened Feb 10, 2025 by
jakevdp
Loading…
Handle extended dtypes within
jax.device_put.
#26415
opened Feb 7, 2025 by
copybara-service
bot
Loading…
jax.numpy reductions: avoid upcast of f16 when dtype is specified by user
pull ready
Ready for copybara import and testing
#26403
opened Feb 7, 2025 by
jakevdp
Loading…
[Pallas/Mosaic GPU] Enable progressive lowering for integer addition.
#26394
opened Feb 7, 2025 by
copybara-service
bot
Loading…
[Pallas/Mosaic GPU] Add initial support for warpgroup semantics in lowering.
#26390
opened Feb 7, 2025 by
copybara-service
bot
Loading…
ProTip!
What’s not been updated in a month: updated:<2025-01-11.