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

python_somax-1.0.1/PKG-INFO

@@ -0,0 +1,202 @@
+Metadata-Version: 2.4
+Name: python-somax
+Version: 1.0.1
+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: Paper, https://arxiv.org/abs/2603.25976
+Keywords: jax,optax,optimization,second-order,gauss-newton,hessian,machine-learning
+Classifier: Development Status :: 3 - Alpha
+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
+Requires-Dist: numpy>=1.23
+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"
+Dynamic: license-file
+
+<h1 align="center">Somax</h1>
+
+<p align="center">
+  <img src="assets/somax_logo_mini.png" alt="Somax logo" width="250px"/>
+</p>
+
+<p align="center">
+  Composable Second-Order Optimization for JAX and Optax.
+</p>
+
+<p align="center">
+  A small research-engineering library for curvature-aware training:
+  modular, matrix-free, and explicit about the moving parts.
+</p>
+
+---
+
+Somax is a JAX-native library for building and running second-order optimization methods from explicit components.
+
+Rather than treating an optimizer as a monolithic object, Somax factors a step into swappable pieces:
+- curvature operator
+- solver
+- damping policy
+- optional preconditioner
+- update transform
+- optional telemetry and control signals
+
+That decomposition is the point.
+
+Somax is built for users who want a clean second-order stack in JAX without hiding the execution model. 
+It aims to make curvature-aware training easier to inspect, compare, and extend.
+
+
+> 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.
+
+
+
+## Why Somax
+
+- **Composable**: build methods from curvature, solver, damping, preconditioner, and update components.
+- **Optax-native**: computed directions are fed through Optax-style update transforms.
+- **Planned execution**: a method is assembled once, planned once, and then executed as a stable step pipeline.
+- **JAX-first**: intended for `jit`-compiled training loops and explicit control over execution.
+- **Multiple solve lanes**: diagonal, parameter-space, and row-space paths are first-class parts of the stack.
+- **Research-friendly**: easy to inspect, compare, ablate, and extend.
+
+
+
+
+## Installation
+
+Install JAX for your backend first:
+
+- JAX installation guide: https://docs.jax.dev/en/latest/installation.html
+
+Then install Somax:
+
+```bash
+pip install python-somax
+```
+
+For local development:
+
+```bash
+git clone https://github.com/cor3bit/somax.git
+cd somax
+pip install -e ".[dev]"
+```
+
+Optional:
+- install `lineax` only if you want to use CG backends with `backend="lineax"`.
+
+
+
+## Quickstart
+
+```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,
+    lam0=1e-2,
+    tol=1e-4,
+    maxiter=20,
+    learning_rate=1e-1,
+)
+
+state = method.init(params)
+
+@jax.jit
+def train_step(params, state, rng):
+    params, state, info = method.step(params, batch, state, rng)
+    return params, state, info
+
+for step in range(10):
+    params, state, info = train_step(params, state, jax.random.fold_in(rng, step))
+```
+
+
+
+
+## Citation
+
+If Somax is useful in your academic work, please cite:
+
+**Second-Order, First-Class: A Composable Stack for Curvature-Aware Training**  
+Mikalai Korbit and Mario Zanon  
+https://arxiv.org/abs/2603.25976
+
+
+```bibtex
+@article{korbit2026second,
+  title={Second-Order, First-Class: A Composable Stack for Curvature-Aware Training},
+  author={Korbit, Mikalai and Zanon, Mario},
+  journal={arXiv preprint arXiv:2603.25976},
+  year={2026}
+}
+```
+
+
+## Related projects
+
+**Optimization in JAX**  
+[Optax](https://github.com/google-deepmind/optax): first-order gradient (e.g., SGD, Adam) optimisers.  
+[JAXopt](https://github.com/google/jaxopt): deterministic second-order methods (e.g., Gauss-Newton, Levenberg
+Marquardt), stochastic first-order methods PolyakSGD, ArmijoSGD.
+
+**Awesome Projects**  
+[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of various JAX projects.  
+[Awesome SOMs](https://github.com/cor3bit/awesome-soms): a list
+of resources for second-order optimization methods in machine learning.
+
+
+
+
+## License
+
+Apache-2.0

python_somax-1.0.1/README.md

@@ -0,0 +1,162 @@
+<h1 align="center">Somax</h1>
+
+<p align="center">
+  <img src="assets/somax_logo_mini.png" alt="Somax logo" width="250px"/>
+</p>
+
+<p align="center">
+  Composable Second-Order Optimization for JAX and Optax.
+</p>
+
+<p align="center">
+  A small research-engineering library for curvature-aware training:
+  modular, matrix-free, and explicit about the moving parts.
+</p>
+
+---
+
+Somax is a JAX-native library for building and running second-order optimization methods from explicit components.
+
+Rather than treating an optimizer as a monolithic object, Somax factors a step into swappable pieces:
+- curvature operator
+- solver
+- damping policy
+- optional preconditioner
+- update transform
+- optional telemetry and control signals
+
+That decomposition is the point.
+
+Somax is built for users who want a clean second-order stack in JAX without hiding the execution model. 
+It aims to make curvature-aware training easier to inspect, compare, and extend.
+
+
+> 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.
+
+
+
+## Why Somax
+
+- **Composable**: build methods from curvature, solver, damping, preconditioner, and update components.
+- **Optax-native**: computed directions are fed through Optax-style update transforms.
+- **Planned execution**: a method is assembled once, planned once, and then executed as a stable step pipeline.
+- **JAX-first**: intended for `jit`-compiled training loops and explicit control over execution.
+- **Multiple solve lanes**: diagonal, parameter-space, and row-space paths are first-class parts of the stack.
+- **Research-friendly**: easy to inspect, compare, ablate, and extend.
+
+
+
+
+## Installation
+
+Install JAX for your backend first:
+
+- JAX installation guide: https://docs.jax.dev/en/latest/installation.html
+
+Then install Somax:
+
+```bash
+pip install python-somax
+```
+
+For local development:
+
+```bash
+git clone https://github.com/cor3bit/somax.git
+cd somax
+pip install -e ".[dev]"
+```
+
+Optional:
+- install `lineax` only if you want to use CG backends with `backend="lineax"`.
+
+
+
+## Quickstart
+
+```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,
+    lam0=1e-2,
+    tol=1e-4,
+    maxiter=20,
+    learning_rate=1e-1,
+)
+
+state = method.init(params)
+
+@jax.jit
+def train_step(params, state, rng):
+    params, state, info = method.step(params, batch, state, rng)
+    return params, state, info
+
+for step in range(10):
+    params, state, info = train_step(params, state, jax.random.fold_in(rng, step))
+```
+
+
+
+
+## Citation
+
+If Somax is useful in your academic work, please cite:
+
+**Second-Order, First-Class: A Composable Stack for Curvature-Aware Training**  
+Mikalai Korbit and Mario Zanon  
+https://arxiv.org/abs/2603.25976
+
+
+```bibtex
+@article{korbit2026second,
+  title={Second-Order, First-Class: A Composable Stack for Curvature-Aware Training},
+  author={Korbit, Mikalai and Zanon, Mario},
+  journal={arXiv preprint arXiv:2603.25976},
+  year={2026}
+}
+```
+
+
+## Related projects
+
+**Optimization in JAX**  
+[Optax](https://github.com/google-deepmind/optax): first-order gradient (e.g., SGD, Adam) optimisers.  
+[JAXopt](https://github.com/google/jaxopt): deterministic second-order methods (e.g., Gauss-Newton, Levenberg
+Marquardt), stochastic first-order methods PolyakSGD, ArmijoSGD.
+
+**Awesome Projects**  
+[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of various JAX projects.  
+[Awesome SOMs](https://github.com/cor3bit/awesome-soms): a list
+of resources for second-order optimization methods in machine learning.
+
+
+
+
+## License
+
+Apache-2.0

python_somax-1.0.1/pyproject.toml

@@ -0,0 +1,93 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "python-somax"
+version = "1.0.1"
+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",
+    "numpy>=1.23",
+]
+keywords = [
+    "jax",
+    "optax",
+    "optimization",
+    "second-order",
+    "gauss-newton",
+    "hessian",
+    "machine-learning",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "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",
+]
+
+[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.0.1/src/python_somax.egg-info/PKG-INFO

@@ -0,0 +1,202 @@
+Metadata-Version: 2.4
+Name: python-somax
+Version: 1.0.1
+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: Paper, https://arxiv.org/abs/2603.25976
+Keywords: jax,optax,optimization,second-order,gauss-newton,hessian,machine-learning
+Classifier: Development Status :: 3 - Alpha
+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
+Requires-Dist: numpy>=1.23
+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"
+Dynamic: license-file
+
+<h1 align="center">Somax</h1>
+
+<p align="center">
+  <img src="assets/somax_logo_mini.png" alt="Somax logo" width="250px"/>
+</p>
+
+<p align="center">
+  Composable Second-Order Optimization for JAX and Optax.
+</p>
+
+<p align="center">
+  A small research-engineering library for curvature-aware training:
+  modular, matrix-free, and explicit about the moving parts.
+</p>
+
+---
+
+Somax is a JAX-native library for building and running second-order optimization methods from explicit components.
+
+Rather than treating an optimizer as a monolithic object, Somax factors a step into swappable pieces:
+- curvature operator
+- solver
+- damping policy
+- optional preconditioner
+- update transform
+- optional telemetry and control signals
+
+That decomposition is the point.
+
+Somax is built for users who want a clean second-order stack in JAX without hiding the execution model. 
+It aims to make curvature-aware training easier to inspect, compare, and extend.
+
+
+> 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.
+
+
+
+## Why Somax
+
+- **Composable**: build methods from curvature, solver, damping, preconditioner, and update components.
+- **Optax-native**: computed directions are fed through Optax-style update transforms.
+- **Planned execution**: a method is assembled once, planned once, and then executed as a stable step pipeline.
+- **JAX-first**: intended for `jit`-compiled training loops and explicit control over execution.
+- **Multiple solve lanes**: diagonal, parameter-space, and row-space paths are first-class parts of the stack.
+- **Research-friendly**: easy to inspect, compare, ablate, and extend.
+
+
+
+
+## Installation
+
+Install JAX for your backend first:
+
+- JAX installation guide: https://docs.jax.dev/en/latest/installation.html
+
+Then install Somax:
+
+```bash
+pip install python-somax
+```
+
+For local development:
+
+```bash
+git clone https://github.com/cor3bit/somax.git
+cd somax
+pip install -e ".[dev]"
+```
+
+Optional:
+- install `lineax` only if you want to use CG backends with `backend="lineax"`.
+
+
+
+## Quickstart
+
+```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,
+    lam0=1e-2,
+    tol=1e-4,
+    maxiter=20,
+    learning_rate=1e-1,
+)
+
+state = method.init(params)
+
+@jax.jit
+def train_step(params, state, rng):
+    params, state, info = method.step(params, batch, state, rng)
+    return params, state, info
+
+for step in range(10):
+    params, state, info = train_step(params, state, jax.random.fold_in(rng, step))
+```
+
+
+
+
+## Citation
+
+If Somax is useful in your academic work, please cite:
+
+**Second-Order, First-Class: A Composable Stack for Curvature-Aware Training**  
+Mikalai Korbit and Mario Zanon  
+https://arxiv.org/abs/2603.25976
+
+
+```bibtex
+@article{korbit2026second,
+  title={Second-Order, First-Class: A Composable Stack for Curvature-Aware Training},
+  author={Korbit, Mikalai and Zanon, Mario},
+  journal={arXiv preprint arXiv:2603.25976},
+  year={2026}
+}
+```
+
+
+## Related projects
+
+**Optimization in JAX**  
+[Optax](https://github.com/google-deepmind/optax): first-order gradient (e.g., SGD, Adam) optimisers.  
+[JAXopt](https://github.com/google/jaxopt): deterministic second-order methods (e.g., Gauss-Newton, Levenberg
+Marquardt), stochastic first-order methods PolyakSGD, ArmijoSGD.
+
+**Awesome Projects**  
+[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of various JAX projects.  
+[Awesome SOMs](https://github.com/cor3bit/awesome-soms): a list
+of resources for second-order optimization methods in machine learning.
+
+
+
+
+## License
+
+Apache-2.0

python_somax-1.0.1/src/python_somax.egg-info/SOURCES.txt

@@ -0,0 +1,53 @@
+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/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/trust_region.py
+src/somax/estimators/__init__.py
+src/somax/estimators/base.py
+src/somax/estimators/gnb_ce.py
+src/somax/estimators/hutchinson.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/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_cholesky.py
+src/somax/solvers/row_common.py