optax

optax copied to clipboard

feat: add `scale_by_gradient_norm`

3

Reference: #594 This PR aims to add a new `GradientTransformation` which allows to scale by the gradient norm. Request for Review: @fabianp

SauravMaheshkar

Add ACProp optimizer.

7

#965

carlosgmartin

**Feature request:** Add the [**ACProp**](https://juntang-zhuang.github.io/acprop/index.html) optimizer. I've created a PR for this: #966.

carlosgmartin

move clipping transforms to optax.transforms.

copybara-service[bot]

"masked" leads to unexpected behavior with callable Pytrees

6

Good morning! The "masked" function finds the "mask_tree" as follows: mask_tree = mask(params) if callable(mask) else mask (1) Which is used twice (in "init_fn" and "update_fn"). However, in some cases,...

JadM133

RST formatting in mechanic

copybara-service[bot]

ntxent fix

8

I added an epsilon value to the cosine similarity function to avoid the NaNs that were occurring when when you had a label vector [0, 0, 0] or when one...

GrantMcConachie

New Example on creating a custom `GradientTransformation`

2

### Motivation Recently for a personal project (#641) I wanted to created a custom `GradientTransformation` object and found out that there wasn't a tutorial or walkthrough on how to create...

SauravMaheshkar

Remove useless inner jit

copybara-service[bot]

**Feature request:** Add a GPU/TPU-friendly solver for the [assignment problem](https://en.wikipedia.org/wiki/Assignment_problem). For context, see: 1. [scipy.optimize.linear_sum_assignment](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.linear_sum_assignment.html) 2. https://github.com/google/jax/issues/10403 3. https://github.com/google/jax/pull/16974 The last page contains the following comment: > There is a...

carlosgmartin