opttx 0.1.0a1__tar.gz

opttx-0.1.0a1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Tianshi Xu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

opttx-0.1.0a1/MANIFEST.in

@@ -0,0 +1,10 @@
+include LICENSE
+include README.md
+include pyproject.toml
+recursive-include opttx *.py
+recursive-include opttx py.typed
+prune ref
+prune tests
+prune docs
+prune examples
+prune .pytest_cache

opttx-0.1.0a1/PKG-INFO

@@ -0,0 +1,220 @@
+Metadata-Version: 2.4
+Name: opttx
+Version: 0.1.0a1
+Summary: JAX/Flax/Optax optimizer manager
+Author: Tianshi Xu
+License: MIT License
+        
+        Copyright (c) 2025 Tianshi Xu
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: Homepage, https://github.com/Hitenze/opttx
+Project-URL: Repository, https://github.com/Hitenze/opttx
+Project-URL: Issues, https://github.com/Hitenze/opttx/issues
+Keywords: jax,optax,flax,optimizer
+Classifier: Development Status :: 3 - Alpha
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.4.20
+Requires-Dist: jaxlib>=0.4.20
+Requires-Dist: optax>=0.2.0
+Requires-Dist: flax>=0.8.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.4; extra == "dev"
+Dynamic: license-file
+
+# OptTx
+
+> **Research Code**: Co-developed with Claude Code, Gemini CLI, Codex CLI, and Cursor. No guarantees provided. Use at your own risk.
+
+JAX/Flax/Optax optimizer library for PINNs and second-order methods.
+
+## Features
+
+- **Multi-term objectives**: `Objective` with `TermSpec` for PINNs (PDE, BC, IC terms)
+- **First-order optimizers**: Adam, SGD, AdamW, SOAP, MUON, Shampoo, L-BFGS
+- **Second-order optimizers**: CGOptimizer (Fisher/GGN), CROptimizer (Hessian)
+- **Acceleration methods**: TGS, NLTGCR, Anderson Acceleration (AA)
+- **Graph neural networks**: GCN, GAT layers for node classification
+- **Matrix-free curvature**: `build_hessian_matvec`, `build_fisher_matvec`
+- **JIT-stable**: Works with `jax.jit` and `jax.lax.scan`
+
+## Install
+
+```bash
+pip install opttx
+```
+
+For development:
+```bash
+pip install -e .[dev]
+```
+
+## Quickstart
+
+### First-order optimizer
+
+```python
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+
+from opttx import Adam, Objective, TermSpec, TrainState
+
+# Define model
+class MLP(nn.Module):
+    @nn.compact
+    def __call__(self, x):
+        x = nn.Dense(32)(x)
+        x = nn.relu(x)
+        x = nn.Dense(1)(x)
+        return x
+
+# Define loss
+def mse_loss(pred, batch):
+    x, y = batch
+    return jnp.mean((pred - y) ** 2)
+
+# Create objective
+term = TermSpec(name="mse", batch_key="data", loss_fn=mse_loss)
+objective = Objective(terms=[term])
+
+# Initialize
+model = MLP()
+params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 3)))["params"]
+
+state = TrainState(
+    step=jnp.array(0),
+    params=params,
+    opt_state=None,
+    apply_fn=lambda v, b: model.apply({"params": v["params"]}, b[0]),
+)
+
+# Create optimizer and train
+optimizer = Adam(objective, learning_rate=1e-3)
+state = optimizer.init(state)
+
+batch = {"data": (jnp.ones((8, 3)), jnp.zeros((8, 1)))}
+state, metrics = optimizer.step(state, batch)
+print(f"Loss: {metrics['loss']}")
+```
+
+### Second-order optimizer (CR + Hessian)
+
+```python
+from opttx import CROptimizer
+
+optimizer = CROptimizer(
+    objective,
+    learning_rate=1.0,
+    damping=1e-3,
+    cr_iters=10,
+    curvature_type="hessian",  # or "fisher"
+)
+state = optimizer.init(state)
+state, metrics = optimizer.step(state, batch)
+```
+
+### Multi-term objective (PINNs)
+
+```python
+def pde_loss(pred, batch):
+    return jnp.mean(pred ** 2)
+
+def bc_loss(pred, batch):
+    return jnp.mean(pred ** 2)
+
+pde_term = TermSpec(name="pde", batch_key="x_pde", loss_fn=pde_loss)
+bc_term = TermSpec(name="bc", batch_key="x_bc", loss_fn=bc_loss)
+
+objective = Objective(
+    terms=[pde_term, bc_term],
+    loss_weights={"pde": 1.0, "bc": 0.1},
+)
+
+batch = {
+    "x_pde": jnp.ones((100, 2)),
+    "x_bc": jnp.ones((20, 2)),
+}
+```
+
+## API Reference
+
+### Optimizers
+
+| Optimizer | Description |
+|-----------|-------------|
+| `Adam` | Adam optimizer |
+| `SGD` | SGD with momentum |
+| `AdamW` | Adam with weight decay |
+| `SOAP` | Second-order approximation |
+| `MUON` | Momentum with orthogonalization |
+| `Shampoo` | Shampoo preconditioner |
+| `LBFGSOptimizer` | L-BFGS quasi-Newton |
+| `CGOptimizer` | Conjugate Gradient (Fisher/GGN) |
+| `CROptimizer` | Conjugate Residual (Hessian) |
+| `TGSOptimizer` | TGS acceleration |
+| `TGSAccelerator` | TGS wrapper for any optimizer |
+| `AAAccelerator` | Anderson Acceleration wrapper |
+| `NLTGCROptimizer` | Nonlinear truncated GCR |
+
+### Curvature
+
+| Function | Description |
+|----------|-------------|
+| `build_hessian_matvec` | Matrix-free Hessian-vector product |
+| `build_fisher_matvec` | Matrix-free Fisher/GGN-vector product |
+| `build_damped_matvec` | Add damping: (H + λI)v |
+
+### Solvers
+
+| Function | Description |
+|----------|-------------|
+| `cg_solve` | Conjugate Gradient solver |
+| `cr_solve` | Conjugate Residual solver |
+| `tgs_solve_fori` | TGS solver (JIT-compatible) |
+| `nltgcr_solve_fori` | NLTGCR solver (JIT-compatible) |
+
+### Models
+
+| Model | Description |
+|-------|-------------|
+| `GCN` | Graph Convolutional Network |
+| `GCNLayer` | Single GCN layer |
+| `GAT` | Graph Attention Network |
+| `GATLayer` | Single GAT layer |
+| `normalize_adjacency` | Symmetric adjacency normalization |
+
+## Design Constraints
+
+- `state.step` must be a scalar `jax.Array` (never Python int)
+- Metrics have static string keys and scalar values
+- Must include `"loss"` key in metrics
+- Multi-term + `batch_stats` is not supported
+
+## License
+
+MIT

opttx-0.1.0a1/README.md

@@ -0,0 +1,172 @@
+# OptTx
+
+> **Research Code**: Co-developed with Claude Code, Gemini CLI, Codex CLI, and Cursor. No guarantees provided. Use at your own risk.
+
+JAX/Flax/Optax optimizer library for PINNs and second-order methods.
+
+## Features
+
+- **Multi-term objectives**: `Objective` with `TermSpec` for PINNs (PDE, BC, IC terms)
+- **First-order optimizers**: Adam, SGD, AdamW, SOAP, MUON, Shampoo, L-BFGS
+- **Second-order optimizers**: CGOptimizer (Fisher/GGN), CROptimizer (Hessian)
+- **Acceleration methods**: TGS, NLTGCR, Anderson Acceleration (AA)
+- **Graph neural networks**: GCN, GAT layers for node classification
+- **Matrix-free curvature**: `build_hessian_matvec`, `build_fisher_matvec`
+- **JIT-stable**: Works with `jax.jit` and `jax.lax.scan`
+
+## Install
+
+```bash
+pip install opttx
+```
+
+For development:
+```bash
+pip install -e .[dev]
+```
+
+## Quickstart
+
+### First-order optimizer
+
+```python
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+
+from opttx import Adam, Objective, TermSpec, TrainState
+
+# Define model
+class MLP(nn.Module):
+    @nn.compact
+    def __call__(self, x):
+        x = nn.Dense(32)(x)
+        x = nn.relu(x)
+        x = nn.Dense(1)(x)
+        return x
+
+# Define loss
+def mse_loss(pred, batch):
+    x, y = batch
+    return jnp.mean((pred - y) ** 2)
+
+# Create objective
+term = TermSpec(name="mse", batch_key="data", loss_fn=mse_loss)
+objective = Objective(terms=[term])
+
+# Initialize
+model = MLP()
+params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 3)))["params"]
+
+state = TrainState(
+    step=jnp.array(0),
+    params=params,
+    opt_state=None,
+    apply_fn=lambda v, b: model.apply({"params": v["params"]}, b[0]),
+)
+
+# Create optimizer and train
+optimizer = Adam(objective, learning_rate=1e-3)
+state = optimizer.init(state)
+
+batch = {"data": (jnp.ones((8, 3)), jnp.zeros((8, 1)))}
+state, metrics = optimizer.step(state, batch)
+print(f"Loss: {metrics['loss']}")
+```
+
+### Second-order optimizer (CR + Hessian)
+
+```python
+from opttx import CROptimizer
+
+optimizer = CROptimizer(
+    objective,
+    learning_rate=1.0,
+    damping=1e-3,
+    cr_iters=10,
+    curvature_type="hessian",  # or "fisher"
+)
+state = optimizer.init(state)
+state, metrics = optimizer.step(state, batch)
+```
+
+### Multi-term objective (PINNs)
+
+```python
+def pde_loss(pred, batch):
+    return jnp.mean(pred ** 2)
+
+def bc_loss(pred, batch):
+    return jnp.mean(pred ** 2)
+
+pde_term = TermSpec(name="pde", batch_key="x_pde", loss_fn=pde_loss)
+bc_term = TermSpec(name="bc", batch_key="x_bc", loss_fn=bc_loss)
+
+objective = Objective(
+    terms=[pde_term, bc_term],
+    loss_weights={"pde": 1.0, "bc": 0.1},
+)
+
+batch = {
+    "x_pde": jnp.ones((100, 2)),
+    "x_bc": jnp.ones((20, 2)),
+}
+```
+
+## API Reference
+
+### Optimizers
+
+| Optimizer | Description |
+|-----------|-------------|
+| `Adam` | Adam optimizer |
+| `SGD` | SGD with momentum |
+| `AdamW` | Adam with weight decay |
+| `SOAP` | Second-order approximation |
+| `MUON` | Momentum with orthogonalization |
+| `Shampoo` | Shampoo preconditioner |
+| `LBFGSOptimizer` | L-BFGS quasi-Newton |
+| `CGOptimizer` | Conjugate Gradient (Fisher/GGN) |
+| `CROptimizer` | Conjugate Residual (Hessian) |
+| `TGSOptimizer` | TGS acceleration |
+| `TGSAccelerator` | TGS wrapper for any optimizer |
+| `AAAccelerator` | Anderson Acceleration wrapper |
+| `NLTGCROptimizer` | Nonlinear truncated GCR |
+
+### Curvature
+
+| Function | Description |
+|----------|-------------|
+| `build_hessian_matvec` | Matrix-free Hessian-vector product |
+| `build_fisher_matvec` | Matrix-free Fisher/GGN-vector product |
+| `build_damped_matvec` | Add damping: (H + λI)v |
+
+### Solvers
+
+| Function | Description |
+|----------|-------------|
+| `cg_solve` | Conjugate Gradient solver |
+| `cr_solve` | Conjugate Residual solver |
+| `tgs_solve_fori` | TGS solver (JIT-compatible) |
+| `nltgcr_solve_fori` | NLTGCR solver (JIT-compatible) |
+
+### Models
+
+| Model | Description |
+|-------|-------------|
+| `GCN` | Graph Convolutional Network |
+| `GCNLayer` | Single GCN layer |
+| `GAT` | Graph Attention Network |
+| `GATLayer` | Single GAT layer |
+| `normalize_adjacency` | Symmetric adjacency normalization |
+
+## Design Constraints
+
+- `state.step` must be a scalar `jax.Array` (never Python int)
+- Metrics have static string keys and scalar values
+- Must include `"loss"` key in metrics
+- Multi-term + `batch_stats` is not supported
+
+## License
+
+MIT

opttx-0.1.0a1/opttx/__init__.py

@@ -0,0 +1,65 @@
+"""OptTx: JAX/Flax/Optax optimizer library for PINNs and second-order methods."""
+
+from .objective import Objective
+from .optimizers import (
+    OptaxOptimizer,
+    Adam,
+    SGD,
+    AdamW,
+    SOAP,
+    Shampoo,
+    MUON,
+    LBFGSOptimizer,
+    CGOptimizer,
+    CROptimizer,
+    NLTGCROptimizer,
+    CrossBatchNLTGCROptimizer,
+    TGSOptimizer,
+    TGSAccelerator,
+    AAAccelerator,
+)
+from .state import TrainState
+from .terms import TermSpec
+from .curvature import build_hessian_matvec, build_fisher_matvec, build_damped_matvec
+from .solvers.cg import cg_solve
+from .solvers.cr import cr_solve
+from .solvers.tgs import tgs_solve_fori
+from .solvers.nltgcr import nltgcr_solve_fori
+
+__version__ = "0.1.0a1"
+
+__all__ = [
+    # First-order optimizers
+    "OptaxOptimizer",
+    "Adam",
+    "SGD",
+    "AdamW",
+    "SOAP",
+    "Shampoo",
+    "MUON",
+    # Quasi-Newton optimizers
+    "LBFGSOptimizer",
+    # Second-order optimizers
+    "CGOptimizer",
+    "CROptimizer",
+    "NLTGCROptimizer",
+    "CrossBatchNLTGCROptimizer",
+    # First-order accelerated
+    "TGSOptimizer",
+    "TGSAccelerator",
+    "AAAccelerator",
+    # Curvature matvecs
+    "build_hessian_matvec",
+    "build_fisher_matvec",
+    "build_damped_matvec",
+    # Solvers
+    "cg_solve",
+    "cr_solve",
+    "tgs_solve_fori",
+    "nltgcr_solve_fori",
+    # Core
+    "Objective",
+    "TermSpec",
+    "TrainState",
+    "__version__",
+]

opttx-0.1.0a1/opttx/apply.py

@@ -0,0 +1,80 @@
+"""Apply wrapper for Flax modules with method dispatch."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, Optional, Sequence, Union
+
+
+def apply_with_method(
+    apply_fn: Callable[..., Any],
+    variables: Dict[str, Any],
+    *args: Any,
+    method: Optional[Union[Callable, str]] = None,
+    training: Optional[bool] = None,
+    rngs: Optional[Dict[str, Any]] = None,
+    mutable: Union[bool, Sequence[str]] = False,
+    **kwargs: Any,
+) -> Any:
+    """Apply a Flax module with optional method dispatch.
+
+    This wrapper enables using different forward methods for multi-term objectives,
+    which is essential for PINNs where different terms may compute different outputs
+    (e.g., residuals, boundary values, initial conditions).
+
+    Args:
+        apply_fn: The Flax module's apply function
+        variables: Model variables dict (must contain "params")
+        *args: Positional arguments to pass to the method
+        method: Optional method to call. Can be:
+            - None: Use default __call__ method
+            - str: Name of module method (e.g., "pde_residual")
+            - Callable: Custom callable(module, *args) -> output
+        training: Whether in training mode. Only passed if explicitly set.
+        rngs: Optional RNG keys dict for stochastic layers
+        mutable: Collections to mark as mutable (e.g., ["batch_stats"])
+        **kwargs: Additional keyword arguments
+
+    Returns:
+        If mutable=False: model output
+        If mutable=True/list: (model_output, updates_dict)
+
+    Example:
+        >>> # Default forward (PINNs - no training flag needed)
+        >>> out = apply_with_method(model.apply, variables, x)
+        >>>
+        >>> # Custom method for PDE residual
+        >>> residual = apply_with_method(
+        ...     model.apply, variables, x, method="compute_residual"
+        ... )
+        >>>
+        >>> # CNN with BatchNorm (training flag needed)
+        >>> out, updates = apply_with_method(
+        ...     model.apply, variables, x,
+        ...     training=True, mutable=["batch_stats"]
+        ... )
+    """
+    # Build apply kwargs - only include training if explicitly set
+    apply_kwargs = {}
+    if training is not None:
+        apply_kwargs["training"] = training
+    apply_kwargs.update(kwargs)
+
+    if rngs is not None:
+        apply_kwargs["rngs"] = rngs
+
+    if mutable:
+        apply_kwargs["mutable"] = mutable
+
+    # Dispatch based on method type
+    if method is None:
+        # Default __call__
+        return apply_fn(variables, *args, **apply_kwargs)
+    elif isinstance(method, str):
+        # String method name - use Flax's method= argument
+        return apply_fn(variables, *args, method=method, **apply_kwargs)
+    elif callable(method):
+        # Callable method - pass module as first arg
+        # Note: Flax apply with method=callable expects the callable to receive module
+        return apply_fn(variables, *args, method=method, **apply_kwargs)
+    else:
+        raise ValueError(f"method must be None, str, or callable, got {type(method)}")