Kei Ishikawa kstoneriv3

kstoneriv3 / optimizers.py

Last active June 25, 2025 02:34

An Optax implementation of Grams, c-Adamw, c-Lion optimizers

	import functools
	import warnings
	from collections.abc import Callable
	from typing import Any, NamedTuple, Optional, Union

	import chex
	import jax
	import jax.numpy as jnp
	import optax.tree
	from jax import nn