Search results

Update JVP rule for abs

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 26086

#26086 In jax-ml/jax;

· dfm opened on Jan 24, 2025

[better_errors] Add debug info to more Jaxprs

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 26078

#26078 In jax-ml/jax;

· gnecula opened on Jan 24, 2025

[better_errors] Port the Pallas debug info mechanisms to the new JAX debugging info

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 26058

#26058 In jax-ml/jax;

· gnecula opened on Jan 23, 2025

Implement lax.full_like in terms of a primitive

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 26045

#26045 In jax-ml/jax;

· jakevdp opened on Jan 22, 2025

Remove linear_util.cache

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 26038

#26038 In jax-ml/jax;

· dougalm opened on Jan 22, 2025

[better_errors] Consolidate TracingDebugInfo and JaxprDebugInfo

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 26001

#26001 In jax-ml/jax;

· gnecula opened on Jan 21, 2025

Clarify type hints for wrapped functions

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 25994

#25994 In jax-ml/jax;

· ncoghlan opened on Jan 21, 2025

[shard_map] fix jax.debug.print/callback under partial auto

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 25887

#25887 In jax-ml/jax;

· mattjj opened on Jan 14, 2025

Add cuDNN flash attention sequence packing

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 25812

#25812 In jax-ml/jax;

· Cjkkkk opened on Jan 9, 2025

Add jax.random.multinomial.

kokoro:force-run

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 25688

#25688 In jax-ml/jax;

· carlosgmartin opened on Dec 27, 2024

Fix debug_nans regressions.

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 25519

#25519 In jax-ml/jax;

· emilyfertig opened on Dec 16, 2024

Simplify implementation of nn.relu.

pull readyReady for copybara import and testingReady for copybara import and testing

Status: Open (in progress).

jax-ml/jaxnumber 25331

#25331 In jax-ml/jax;

· carlosgmartin opened on Dec 8, 2024