python-somax 0.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.1.0/pyproject.toml

@@ -0,0 +1,94 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "python-somax"
+version = "1.1.0"
+description = "Composable Second-Order Optimization for JAX and Optax."
+readme = "README.md"
+requires-python = ">=3.11"
+license = "Apache-2.0"
+license-files = ["LICENSE"]
+authors = [
+    { name = "Nick Korbit" }
+]
+dependencies = [
+    "jax>=0.4.20",
+    "optax>=0.1.7",
+    "flax>=0.7.5",
+]
+keywords = [
+    "jax",
+    "optax",
+    "optimization",
+    "second-order",
+    "gauss-newton",
+    "hessian",
+    "machine-learning",
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Mathematics",
+]
+
+[project.urls]
+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"
+Paper = "https://arxiv.org/abs/2603.25976"
+
+[project.optional-dependencies]
+dev = [
+    "build>=1.2",
+    "twine>=5",
+    "pytest>=8",
+    "pytest-cov>=5",
+    "ruff>=0.5",
+    "pyright>=1.1",
+    "pre-commit>=3",
+    "chex>=0.1.8",
+]
+docs = [
+    "mkdocs>=1.5",
+    "mkdocs-material>=9",
+]
+matfree = ["matfree"]
+lineax = ["lineax"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "N", "W", "UP", "B", "SIM"]
+ignore = ["E501", "N803", "N806"]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "lf"
+
+[tool.pyright]
+include = ["src"]
+exclude = ["**/node_modules", "**/__pycache__"]
+reportMissingImports = true
+reportMissingTypeStubs = false
+pythonVersion = "3.11"
+typeCheckingMode = "basic"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-q --strict-markers"
+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.1.0/src/python_somax.egg-info/SOURCES.txt

@@ -0,0 +1,58 @@
+LICENSE
+README.md
+pyproject.toml
+src/python_somax.egg-info/PKG-INFO
+src/python_somax.egg-info/SOURCES.txt
+src/python_somax.egg-info/dependency_links.txt
+src/python_somax.egg-info/requires.txt
+src/python_somax.egg-info/top_level.txt
+src/somax/__init__.py
+src/somax/assembler.py
+src/somax/executor.py
+src/somax/metrics.py
+src/somax/optax.py
+src/somax/planner.py
+src/somax/presets.py
+src/somax/specs.py
+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
+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
+src/somax/methods/sophia_h.py
+src/somax/preconditioners/__init__.py
+src/somax/preconditioners/base.py
+src/somax/preconditioners/diag_direct.py
+src/somax/preconditioners/diag_ema.py
+src/somax/preconditioners/identity.py
+src/somax/solvers/__init__.py
+src/somax/solvers/base.py
+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_cg_lifted.py
+src/somax/solvers/row_common.py

python_somax-1.1.0/src/python_somax.egg-info/requires.txt

@@ -0,0 +1,23 @@
+jax>=0.4.20
+optax>=0.1.7
+flax>=0.7.5
+
+[dev]
+build>=1.2
+twine>=5
+pytest>=8
+pytest-cov>=5
+ruff>=0.5
+pyright>=1.1
+pre-commit>=3
+chex>=0.1.8
+
+[docs]
+mkdocs>=1.5
+mkdocs-material>=9
+
+[lineax]
+lineax
+
+[matfree]
+matfree