Search results

Extend Muon optimizer to support callable masks in weight_dimension_numbers

type:featureNew feature or requestNew feature or request

Status: Open.

#1537 In google-deepmind/optax;

· yash194 opened on Dec 24, 2025

Add optional fast path for <code>clip_by_global_norm</code> when differentiability is not required

type:featureNew feature or requestNew feature or request

Status: Open.

#1535 In google-deepmind/optax;

· SanchitKS12 opened on Dec 17, 2025

<code>inject_hyperparams</code> causes numerical divergence in Adam compared to standard optimizer with identical schedule

type:bugSomething isn't workingSomething isn't working

Status: Open.

#1532 In google-deepmind/optax;

· kameronton opened on Dec 11, 2025

Feature request: add BYOL, SimSiam, DINO and Barlow Twins losses to <code>optax.losses</code>

type:featureNew feature or requestNew feature or request

Status: Open.

#1528 In google-deepmind/optax;

· surajyadav-research opened on Dec 10, 2025

Code for DoG is duplicated

type:others

Status: Open.

#1509 In google-deepmind/optax;

· carlosgmartin opened on Nov 23, 2025

Feature Request: Auction algorithm for linear assignment

type:featureNew feature or requestNew feature or request

Status: Open.

#1507 In google-deepmind/optax;

· surajyadav-research opened on Nov 20, 2025

Should <code>tree_sum</code> use <code>jax.tree.reduce_associative</code>?

type:supportFurther information is requestedFurther information is requested

Status: Open.

#1498 In google-deepmind/optax;

· SobhanMP opened on Nov 17, 2025

Add <code>label_smoothing</code> argument to <code>softmax_cross_entropy_with_integer_labels</code>

type:featureNew feature or requestNew feature or request

Status: Open.

#1486 In google-deepmind/optax;

· carlosgmartin opened on Nov 12, 2025

Add Optional Fast Path for <code>clip_by_global_norm</code> (opt-in non-differentiable mode)

type:featureNew feature or requestNew feature or request

Status: Open.

#1452 In google-deepmind/optax;

· satyam-mishra-dev opened on Oct 29, 2025

<code>reduce_on_plateau</code> causes JIT cache miss due to weak-typed <code>avg_value</code>

type:bugSomething isn't workingSomething isn't working

Status: Open.

#1440 In google-deepmind/optax;

· Mycroft-47 opened on Oct 24, 2025

Add Implementation of _Lion Optimizer (Evolved sign Momentum)

type:featureNew feature or requestNew feature or request

Status: Open.

#1437 In google-deepmind/optax;

· raghulchandramouli opened on Oct 24, 2025

Correct usage of look-ahead optimizer

type:supportFurther information is requestedFurther information is requested

Status: Open.

#1429 In google-deepmind/optax;

· SNMS95 opened on Oct 7, 2025