python-somax 1.0.1__tar.gz → 1.1.0__tar.gz

python_somax-1.1.0/PKG-INFO

@@ -0,0 +1,182 @@
+Metadata-Version: 2.4
+Name: python-somax
+Version: 1.1.0
+Summary: Composable Second-Order Optimization for JAX and Optax.
+Author: Nick Korbit
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/cor3bit/somax
+Project-URL: Repository, https://github.com/cor3bit/somax
+Project-URL: Issues, https://github.com/cor3bit/somax/issues
+Project-URL: Releases, https://github.com/cor3bit/somax/releases
+Project-URL: Documentation, https://github.com/cor3bit/somax/tree/main/docs
+Project-URL: Paper, https://arxiv.org/abs/2603.25976
+Keywords: jax,optax,optimization,second-order,gauss-newton,hessian,machine-learning
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.4.20
+Requires-Dist: optax>=0.1.7
+Requires-Dist: flax>=0.7.5
+Provides-Extra: dev
+Requires-Dist: build>=1.2; extra == "dev"
+Requires-Dist: twine>=5; extra == "dev"
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: pytest-cov>=5; extra == "dev"
+Requires-Dist: ruff>=0.5; extra == "dev"
+Requires-Dist: pyright>=1.1; extra == "dev"
+Requires-Dist: pre-commit>=3; extra == "dev"
+Requires-Dist: chex>=0.1.8; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: mkdocs>=1.5; extra == "docs"
+Requires-Dist: mkdocs-material>=9; extra == "docs"
+Provides-Extra: matfree
+Requires-Dist: matfree; extra == "matfree"
+Provides-Extra: lineax
+Requires-Dist: lineax; extra == "lineax"
+Dynamic: license-file
+
+<h1 align="center">Somax</h1>
+
+<p align="center">
+  <img src="assets/somax_logo_mini.png" alt="Somax logo" width="250" />
+</p>
+
+<p align="center">
+  <strong>Second-Order Methods in JAX</strong>
+</p>
+
+
+
+Somax is a JAX-native library for composable second-order optimization.
+It combines curvature operators, solvers, damping, preconditioners,
+and update transforms into flexible optimization methods.
+
+
+Features:
+- **Composable:** combine curvature, solvers, damping, preconditioners, and update rules.
+- **Matrix-free:** compute curvature products without forming dense matrices.
+- **JAX-first:** JIT-compiled steps with explicit control over execution.
+- **Optax integration:** apply standard Optax transformations to computed directions.
+- **Research-friendly:** swap components, run ablations, and prototype new methods.
+
+
+
+
+> The catfish in the logo is a small nod to *som*, the Belarusian word for catfish. 
+> A quiet bottom-dweller, but not a first-order creature.
+
+
+
+
+
+
+## Installation
+
+Install [JAX](https://docs.jax.dev/en/latest/installation.html) for your hardware, then:
+
+```bash
+pip install python-somax
+```
+
+For development:
+
+```bash
+git clone https://github.com/cor3bit/somax.git
+cd somax
+pip install -e ".[dev]"
+```
+
+The `lineax` package is optional and needed only for CG with `backend="lineax"`.
+
+
+
+
+
+
+## Quickstart
+
+
+A minimal stochastic Gauss-Newton training loop for multiclass classification:
+
+```python
+import jax
+import jax.numpy as jnp
+import somax
+
+
+def predict_fn(params, x):
+    h = jnp.tanh(x @ params["W1"] + params["b1"])
+    return h @ params["W2"] + params["b2"]
+
+
+rng = jax.random.PRNGKey(0)
+k1, k2, k3, k4 = jax.random.split(rng, 4)
+
+params = {
+    "W1": 0.1 * jax.random.normal(k1, (16, 32)),
+    "b1": jnp.zeros((32,)),
+    "W2": 0.1 * jax.random.normal(k2, (32, 10)),
+    "b2": jnp.zeros((10,)),
+}
+
+batch = {
+    "x": jax.random.normal(k3, (64, 16)),
+    "y": jax.random.randint(k4, (64,), 0, 10),
+}
+
+method = somax.sgn_ce(
+    predict_fn=predict_fn,
+    learning_rate=0.1,
+    lam0=1.0,
+    maxiter=10,
+    tol=1e-4,
+)
+
+state = method.init(params)
+
+
+@jax.jit
+def train_step(params, batch, state, rng):
+    return method.step(params, batch, state, rng)
+
+
+for step in range(10):
+    step_rng = jax.random.fold_in(rng, step)
+    params, state, info = train_step(params, batch, state, step_rng)
+```
+
+For more examples, see [Usage: A paper in a few lines](docs/USAGE.md).
+
+
+
+
+## Citation
+
+Mikalai Korbit and Mario Zanon,
+"Second-Order, First-Class: A Composable Stack for Curvature-Aware Training",
+ECML PKDD 2026, [arXiv](https://arxiv.org/abs/2603.25976).
+
+```bibtex
+@inproceedings{korbit2026second,
+  title={Second-Order, First-Class: A Composable Stack for Curvature-Aware Training},
+  author={Korbit, Mikalai and Zanon, Mario},
+  booktitle={Joint European Conference on Machine Learning and Knowledge Discovery in Databases},
+  pages={90--105},
+  year={2026},
+  organization={Springer}
+}
+```
+
+
+
+## License
+
+Apache-2.0
+

python_somax-1.1.0/README.md

@@ -0,0 +1,138 @@
+<h1 align="center">Somax</h1>
+
+<p align="center">
+  <img src="assets/somax_logo_mini.png" alt="Somax logo" width="250" />
+</p>
+
+<p align="center">
+  <strong>Second-Order Methods in JAX</strong>
+</p>
+
+
+
+Somax is a JAX-native library for composable second-order optimization.
+It combines curvature operators, solvers, damping, preconditioners,
+and update transforms into flexible optimization methods.
+
+
+Features:
+- **Composable:** combine curvature, solvers, damping, preconditioners, and update rules.
+- **Matrix-free:** compute curvature products without forming dense matrices.
+- **JAX-first:** JIT-compiled steps with explicit control over execution.
+- **Optax integration:** apply standard Optax transformations to computed directions.
+- **Research-friendly:** swap components, run ablations, and prototype new methods.
+
+
+
+
+> The catfish in the logo is a small nod to *som*, the Belarusian word for catfish. 
+> A quiet bottom-dweller, but not a first-order creature.
+
+
+
+
+
+
+## Installation
+
+Install [JAX](https://docs.jax.dev/en/latest/installation.html) for your hardware, then:
+
+```bash
+pip install python-somax
+```
+
+For development:
+
+```bash
+git clone https://github.com/cor3bit/somax.git
+cd somax
+pip install -e ".[dev]"
+```
+
+The `lineax` package is optional and needed only for CG with `backend="lineax"`.
+
+
+
+
+
+
+## Quickstart
+
+
+A minimal stochastic Gauss-Newton training loop for multiclass classification:
+
+```python
+import jax
+import jax.numpy as jnp
+import somax
+
+
+def predict_fn(params, x):
+    h = jnp.tanh(x @ params["W1"] + params["b1"])
+    return h @ params["W2"] + params["b2"]
+
+
+rng = jax.random.PRNGKey(0)
+k1, k2, k3, k4 = jax.random.split(rng, 4)
+
+params = {
+    "W1": 0.1 * jax.random.normal(k1, (16, 32)),
+    "b1": jnp.zeros((32,)),
+    "W2": 0.1 * jax.random.normal(k2, (32, 10)),
+    "b2": jnp.zeros((10,)),
+}
+
+batch = {
+    "x": jax.random.normal(k3, (64, 16)),
+    "y": jax.random.randint(k4, (64,), 0, 10),
+}
+
+method = somax.sgn_ce(
+    predict_fn=predict_fn,
+    learning_rate=0.1,
+    lam0=1.0,
+    maxiter=10,
+    tol=1e-4,
+)
+
+state = method.init(params)
+
+
+@jax.jit
+def train_step(params, batch, state, rng):
+    return method.step(params, batch, state, rng)
+
+
+for step in range(10):
+    step_rng = jax.random.fold_in(rng, step)
+    params, state, info = train_step(params, batch, state, step_rng)
+```
+
+For more examples, see [Usage: A paper in a few lines](docs/USAGE.md).
+
+
+
+
+## Citation
+
+Mikalai Korbit and Mario Zanon,
+"Second-Order, First-Class: A Composable Stack for Curvature-Aware Training",
+ECML PKDD 2026, [arXiv](https://arxiv.org/abs/2603.25976).
+
+```bibtex
+@inproceedings{korbit2026second,
+  title={Second-Order, First-Class: A Composable Stack for Curvature-Aware Training},
+  author={Korbit, Mikalai and Zanon, Mario},
+  booktitle={Joint European Conference on Machine Learning and Knowledge Discovery in Databases},
+  pages={90--105},
+  year={2026},
+  organization={Springer}
+}
+```
+
+
+
+## License
+
+Apache-2.0
+

{python_somax-1.0.1 → python_somax-1.1.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "python-somax"
-version = "1.0.1"
+version = "1.1.0"
 description = "Composable Second-Order Optimization for JAX and Optax."
 readme = "README.md"
 requires-python = ">=3.11"
@@ -17,7 +17,6 @@ dependencies = [
     "jax>=0.4.20",
     "optax>=0.1.7",
     "flax>=0.7.5",
-    "numpy>=1.23",
 ]
 keywords = [
     "jax",
@@ -29,7 +28,7 @@ keywords = [
     "machine-learning",
 ]
 classifiers = [
-    "Development Status :: 3 - Alpha",
+    "Development Status :: 4 - Beta",
     "Intended Audience :: Science/Research",
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.11",
@@ -43,7 +42,7 @@ Homepage = "https://github.com/cor3bit/somax"
 Repository = "https://github.com/cor3bit/somax"
 Issues = "https://github.com/cor3bit/somax/issues"
 Releases = "https://github.com/cor3bit/somax/releases"
-#Documentation = "https://github.com/cor3bit/somax/tree/main/docs"
+Documentation = "https://github.com/cor3bit/somax/tree/main/docs"
 Paper = "https://arxiv.org/abs/2603.25976"
 
 [project.optional-dependencies]
@@ -61,6 +60,8 @@ docs = [
     "mkdocs>=1.5",
     "mkdocs-material>=9",
 ]
+matfree = ["matfree"]
+lineax = ["lineax"]
 
 [tool.setuptools.packages.find]
 where = ["src"]
@@ -90,4 +91,4 @@ typeCheckingMode = "basic"
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 addopts = "-q --strict-markers"
-xfail_strict = true
+xfail_strict = true

python_somax-1.1.0/src/python_somax.egg-info/PKG-INFO

@@ -0,0 +1,182 @@
+Metadata-Version: 2.4
+Name: python-somax
+Version: 1.1.0
+Summary: Composable Second-Order Optimization for JAX and Optax.
+Author: Nick Korbit
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/cor3bit/somax
+Project-URL: Repository, https://github.com/cor3bit/somax
+Project-URL: Issues, https://github.com/cor3bit/somax/issues
+Project-URL: Releases, https://github.com/cor3bit/somax/releases
+Project-URL: Documentation, https://github.com/cor3bit/somax/tree/main/docs
+Project-URL: Paper, https://arxiv.org/abs/2603.25976
+Keywords: jax,optax,optimization,second-order,gauss-newton,hessian,machine-learning
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.4.20
+Requires-Dist: optax>=0.1.7
+Requires-Dist: flax>=0.7.5
+Provides-Extra: dev
+Requires-Dist: build>=1.2; extra == "dev"
+Requires-Dist: twine>=5; extra == "dev"
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: pytest-cov>=5; extra == "dev"
+Requires-Dist: ruff>=0.5; extra == "dev"
+Requires-Dist: pyright>=1.1; extra == "dev"
+Requires-Dist: pre-commit>=3; extra == "dev"
+Requires-Dist: chex>=0.1.8; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: mkdocs>=1.5; extra == "docs"
+Requires-Dist: mkdocs-material>=9; extra == "docs"
+Provides-Extra: matfree
+Requires-Dist: matfree; extra == "matfree"
+Provides-Extra: lineax
+Requires-Dist: lineax; extra == "lineax"
+Dynamic: license-file
+
+<h1 align="center">Somax</h1>
+
+<p align="center">
+  <img src="assets/somax_logo_mini.png" alt="Somax logo" width="250" />
+</p>
+
+<p align="center">
+  <strong>Second-Order Methods in JAX</strong>
+</p>
+
+
+
+Somax is a JAX-native library for composable second-order optimization.
+It combines curvature operators, solvers, damping, preconditioners,
+and update transforms into flexible optimization methods.
+
+
+Features:
+- **Composable:** combine curvature, solvers, damping, preconditioners, and update rules.
+- **Matrix-free:** compute curvature products without forming dense matrices.
+- **JAX-first:** JIT-compiled steps with explicit control over execution.
+- **Optax integration:** apply standard Optax transformations to computed directions.
+- **Research-friendly:** swap components, run ablations, and prototype new methods.
+
+
+
+
+> The catfish in the logo is a small nod to *som*, the Belarusian word for catfish. 
+> A quiet bottom-dweller, but not a first-order creature.
+
+
+
+
+
+
+## Installation
+
+Install [JAX](https://docs.jax.dev/en/latest/installation.html) for your hardware, then:
+
+```bash
+pip install python-somax
+```
+
+For development:
+
+```bash
+git clone https://github.com/cor3bit/somax.git
+cd somax
+pip install -e ".[dev]"
+```
+
+The `lineax` package is optional and needed only for CG with `backend="lineax"`.
+
+
+
+
+
+
+## Quickstart
+
+
+A minimal stochastic Gauss-Newton training loop for multiclass classification:
+
+```python
+import jax
+import jax.numpy as jnp
+import somax
+
+
+def predict_fn(params, x):
+    h = jnp.tanh(x @ params["W1"] + params["b1"])
+    return h @ params["W2"] + params["b2"]
+
+
+rng = jax.random.PRNGKey(0)
+k1, k2, k3, k4 = jax.random.split(rng, 4)
+
+params = {
+    "W1": 0.1 * jax.random.normal(k1, (16, 32)),
+    "b1": jnp.zeros((32,)),
+    "W2": 0.1 * jax.random.normal(k2, (32, 10)),
+    "b2": jnp.zeros((10,)),
+}
+
+batch = {
+    "x": jax.random.normal(k3, (64, 16)),
+    "y": jax.random.randint(k4, (64,), 0, 10),
+}
+
+method = somax.sgn_ce(
+    predict_fn=predict_fn,
+    learning_rate=0.1,
+    lam0=1.0,
+    maxiter=10,
+    tol=1e-4,
+)
+
+state = method.init(params)
+
+
+@jax.jit
+def train_step(params, batch, state, rng):
+    return method.step(params, batch, state, rng)
+
+
+for step in range(10):
+    step_rng = jax.random.fold_in(rng, step)
+    params, state, info = train_step(params, batch, state, step_rng)
+```
+
+For more examples, see [Usage: A paper in a few lines](docs/USAGE.md).
+
+
+
+
+## Citation
+
+Mikalai Korbit and Mario Zanon,
+"Second-Order, First-Class: A Composable Stack for Curvature-Aware Training",
+ECML PKDD 2026, [arXiv](https://arxiv.org/abs/2603.25976).
+
+```bibtex
+@inproceedings{korbit2026second,
+  title={Second-Order, First-Class: A Composable Stack for Curvature-Aware Training},
+  author={Korbit, Mikalai and Zanon, Mario},
+  booktitle={Joint European Conference on Machine Learning and Knowledge Discovery in Databases},
+  pages={90--105},
+  year={2026},
+  organization={Springer}
+}
+```
+
+
+
+## License
+
+Apache-2.0
+

{python_somax-1.0.1 → python_somax-1.1.0}/src/python_somax.egg-info/SOURCES.txt

@@ -18,6 +18,7 @@ src/somax/types.py
 src/somax/utils.py
 src/somax/curvature/__init__.py
 src/somax/curvature/base.py
+src/somax/curvature/fgnha_ce.py
 src/somax/curvature/ggn_ce.py
 src/somax/curvature/ggn_mse.py
 src/somax/curvature/hessian.py
@@ -25,15 +26,19 @@ src/somax/curvature/with_estimators.py
 src/somax/damping/__init__.py
 src/somax/damping/base.py
 src/somax/damping/constant.py
+src/somax/damping/scheduled.py
 src/somax/damping/trust_region.py
 src/somax/estimators/__init__.py
 src/somax/estimators/base.py
+src/somax/estimators/diagpp.py
 src/somax/estimators/gnb_ce.py
 src/somax/estimators/hutchinson.py
+src/somax/estimators/lanczos.py
 src/somax/methods/__init__.py
 src/somax/methods/adahessian.py
 src/somax/methods/direct_methods.py
 src/somax/methods/egn.py
+src/somax/methods/fgn.py
 src/somax/methods/newton_cg.py
 src/somax/methods/sgn.py
 src/somax/methods/sophia_g.py
@@ -49,5 +54,5 @@ src/somax/solvers/cg.py
 src/somax/solvers/direct.py
 src/somax/solvers/identity.py
 src/somax/solvers/row_cg.py
-src/somax/solvers/row_cholesky.py
+src/somax/solvers/row_cg_lifted.py
 src/somax/solvers/row_common.py

{python_somax-1.0.1 → python_somax-1.1.0}/src/python_somax.egg-info/requires.txt

@@ -1,7 +1,6 @@
 jax>=0.4.20
 optax>=0.1.7
 flax>=0.7.5
-numpy>=1.23
 
 [dev]
 build>=1.2
@@ -16,3 +15,9 @@ chex>=0.1.8
 [docs]
 mkdocs>=1.5
 mkdocs-material>=9
+
+[lineax]
+lineax
+
+[matfree]
+matfree

{python_somax-1.0.1 → python_somax-1.1.0}/src/somax/__init__.py

@@ -20,9 +20,12 @@ from .methods import (
     adahessian,
     ggn_direct_mse,
     ggn_direct_ce,
+    fgn_direct_ce,
     newton_direct,
     egn_ce,
     egn_mse,
+    fgn_ce,
+    fgn_ce_exact,
     newton_cg,
     sgn_mse,
     sgn_ce,
@@ -47,9 +50,12 @@ __all__ = [
     "adahessian",
     "ggn_direct_mse",
     "ggn_direct_ce",
+    "fgn_direct_ce",
     "newton_direct",
     "egn_ce",
     "egn_mse",
+    "fgn_ce",
+    "fgn_ce_exact",
     "newton_cg",
     "sgn_mse",
     "sgn_ce",

{python_somax-1.0.1 → python_somax-1.1.0}/src/somax/assembler.py

@@ -33,6 +33,8 @@ from .types import Array, Batch, Params, PRNGKey, StepInfo, ScalarOrSchedule, Gr
 class SecondOrderState:
     """Single state tree owned by Somax.
 
+    model_state: optional non-parameter model state, such as BatchNorm running
+        statistics. It is updated once by the curvature primal each outer step.
     opt_state: post-direction tx state (Optax state or custom).
     method_state: optional extra state for method-specific algorithms (eg AdaHessian EMAs).
     """
@@ -42,13 +44,14 @@ class SecondOrderState:
     solver_state: Any
     opt_state: Any
     method_state: Any  # optional extension point (eg AdaHessian)
+    model_state: Any = ()
 
 
 class SecondOrderMethod:
-    """Runnable Somax method: init/update.
+    """Runnable Somax method: init/step.
 
-    update(...) performs one fused step:
-      init cstate -> grad/row_op -> solve direction -> tx.update -> apply -> (optional rho) -> damping.update
+    step(...) performs one fused step:
+      init cstate -> grad/row_op -> solve direction -> tx.update -> apply -> (optional gain ratio) -> damping.update
     """
 
     def __init__(
@@ -68,11 +71,18 @@ class SecondOrderMethod:
         self.plan = plan
         self.tx = tx
 
-    def init(self, params: Params) -> SecondOrderState:
+    def init(
+            self,
+            params: Params,
+            *,
+            model_state: Any = (),
+    ) -> SecondOrderState:
+        """Initialize optimizer and optional non-parameter model state."""
         m_t = jax.tree_util.tree_map(jnp.zeros_like, params) if self.plan.want.want_m else ()
 
         return SecondOrderState(
             step=jnp.asarray(0, jnp.int32),
+            model_state=model_state,
             damp_state=self.damping.init(),
             precond_state=self.precond.init(params),
             solver_state=self.solver.init(params),
@@ -118,8 +128,8 @@ def assemble(
         estimator: Optional[EstimatorSpec] = None,
         telemetry: Optional[TelemetrySpec] = None,
 
-        # post-direction application
-        tx: GradientTransformation | None = None,  # full optax chain, if None, we build from knobs below
+        # post-direction application (owned)
+        tx: GradientTransformation | None = None,  # full optax chain (expert); if None, we build from knobs below
         learning_rate: ScalarOrSchedule = 1.0,
         clip_norm: float | None = None,
         weight_decay: float | None = None,
@@ -137,7 +147,7 @@ def assemble(
     # Build plan from components
     plan = build_plan(damping=dp, solver=slv, precond=pp, telemetry_reqs=tlm_reqs)
 
-    # Build owned post-direction tx
+    # Build owned post-direction tx (identity/sign+lr/clip/wd), unless user provided a full tx.
     optax_chain = build_optax_tx(
         tx=tx,
         learning_rate=learning_rate,