zeroth-learn 0.2.4__tar.gz → 0.2.6__tar.gz

zeroth_learn-0.2.6/PKG-INFO

@@ -0,0 +1,92 @@
+Metadata-Version: 2.4
+Name: zeroth-learn
+Version: 0.2.6
+Summary: Zeroth-order optimization for black-box models
+Project-URL: Repository, https://github.com/nicolasmalet/Zeroth-Learn
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: matplotlib
+Requires-Dist: cycler
+Provides-Extra: mnist
+Requires-Dist: scikit-learn; extra == "mnist"
+
+# Zeroth-Learn
+
+**A NumPy library for estimating gradients from loss evaluations and training black-box models.**
+
+The library implements coordinate finite differences, random-direction estimation, and SGD/Adam updates. Its MNIST experiments check that these components can train models when backpropagation is unavailable. The estimator and optimizer interfaces can also be used with other black boxes, including the simulator in [Quantum-Learn](https://github.com/nicolasmalet/Quantum-Learn).
+
+## 1. Optimisation without derivatives
+
+Let $\theta\in\mathbb R^d$ be model parameters and $f(\theta)$ a mini-batch loss. The model exposes values of $f$, but not $\nabla f$. Zeroth-order optimisation estimates a descent direction by evaluating nearby parameters:
+
+```math
+\theta_{k+1}=\theta_k-\eta_k\widehat{\nabla f}(\theta_k).
+```
+
+For Adam, the estimated gradient enters the usual first- and second-moment update in place of an exact gradient. No derivative of the black-box model is required.
+
+## 2. Finite-difference estimators
+
+A one-sided coordinate difference estimates component $i$ as
+
+```math
+\widehat{\partial_i f}(\theta)
+=\frac{f(\theta+\delta e_i)-f(\theta)}{\delta}.
+```
+
+Estimating all $d$ components needs $d+1$ loss evaluations. [`GlobalFiniteDifference`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/gradient_estimators.py) does this for every coordinate; [`PartialFiniteDifference`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/gradient_estimators.py) evaluates only selected coordinates and returns zero elsewhere.
+
+For a cheaper estimate in high dimension, [`SimultaneousPerturbation`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/gradient_estimators.py) draws $T$ Rademacher directions. At construction, the matrix $P\in\mathbb R^{T\times d}$ has entries $P_{ij}=s_{ij}/\sqrt T$, where each $s_{ij}$ is $+1$ or $-1$ with equal probability. The implementation uses
+
+```math
+\widehat{\nabla f}(\theta)
+=\frac{1}{\delta}P^\top
+\begin{pmatrix}
+f(\theta+\delta P_1)-f(\theta)\\
+\vdots\\
+f(\theta+\delta P_T)-f(\theta)
+\end{pmatrix}.
+```
+
+Because $\mathbb E[P^\top P]=I$, the first-order term recovers $\nabla f$ in expectation. Each update costs $T+1$ loss evaluations, independent of $d$ for a fixed $T$. Larger $T$ averages more directions but increases evaluation cost and memory.
+
+## 3. From an estimator to a model update
+
+The interfaces separate three operations:
+
+| Operation                                                  | Implementation                                                              |
+| ---------------------------------------------------------- | --------------------------------------------------------------------------- |
+| Generate perturbed parameters and reconstruct the gradient | [`gradient_estimators.py`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/gradient_estimators.py)     |
+| Generate Rademacher directions                             | [`perturbation_matrices.py`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/utils/perturbation_matrices.py)        |
+| Evaluate a model at nominal and perturbed parameters       | [`zeroth_order_blackbox.py`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/zeroth_order_blackbox.py) |
+| Apply SGD or Adam to the estimated gradient                | [`optimizers.py`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/optimizers.py)                       |
+
+The included [`ZerothOrderNeuralNetwork`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/neural_network/neural_network.py) is one black-box implementation. [`ParameterManager`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/neural_network/parameter_manager.py) maps weights and biases to a flat $\theta$; NumPy broadcasting evaluates its $T+1$ perturbed networks in one batch. This batching reduces Python overhead, not the number of loss evaluations. Other models choose how to perform those evaluations.
+
+## 4. Numerical checks
+
+A deterministic [test](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/tests/test_gradient_estimator.py) applies the random-direction estimator to a linear function, whose gradient is known. The MNIST experiments provide end-to-end checks of zeroth-order training on CPU. Their protocols, results, plots, and raw loss traces are in [Experiments](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/EXPERIMENTS.md).
+
+## 5. Reproduce the experiments
+
+With [uv](https://docs.astral.sh/uv/), the lockfile fixes the dependency versions. The first experiment run downloads MNIST from OpenML.
+
+```bash
+git clone https://github.com/nicolasmalet/Zeroth-Learn.git
+cd Zeroth-Learn
+uv sync --locked --extra mnist
+uv run --locked --extra mnist python -m unittest discover -s tests
+MPLBACKEND=Agg uv run --locked --extra mnist python -m lab.mnist
+MPLBACKEND=Agg uv run --locked --extra mnist python -m lab.mnist.run_experiment nb_perturbations_vs_model_size
+```
+
+The experiments write accuracy, loss traces, and plots under `results/`.
+
+Without uv, install the package with its MNIST extra in a virtual environment, then run the same Python modules.
+
+## 6. Scope
+
+The MNIST experiments validate classical models. They are not benchmarks of quantum hardware or of PyTorch, and their saved plots do not contain repeated trials or uncertainty estimates. Memory use grows with the number of perturbations in the vectorized neural-network implementation.

zeroth_learn-0.2.6/README.md

@@ -0,0 +1,78 @@
+# Zeroth-Learn
+
+**A NumPy library for estimating gradients from loss evaluations and training black-box models.**
+
+The library implements coordinate finite differences, random-direction estimation, and SGD/Adam updates. Its MNIST experiments check that these components can train models when backpropagation is unavailable. The estimator and optimizer interfaces can also be used with other black boxes, including the simulator in [Quantum-Learn](https://github.com/nicolasmalet/Quantum-Learn).
+
+## 1. Optimisation without derivatives
+
+Let $\theta\in\mathbb R^d$ be model parameters and $f(\theta)$ a mini-batch loss. The model exposes values of $f$, but not $\nabla f$. Zeroth-order optimisation estimates a descent direction by evaluating nearby parameters:
+
+```math
+\theta_{k+1}=\theta_k-\eta_k\widehat{\nabla f}(\theta_k).
+```
+
+For Adam, the estimated gradient enters the usual first- and second-moment update in place of an exact gradient. No derivative of the black-box model is required.
+
+## 2. Finite-difference estimators
+
+A one-sided coordinate difference estimates component $i$ as
+
+```math
+\widehat{\partial_i f}(\theta)
+=\frac{f(\theta+\delta e_i)-f(\theta)}{\delta}.
+```
+
+Estimating all $d$ components needs $d+1$ loss evaluations. [`GlobalFiniteDifference`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/gradient_estimators.py) does this for every coordinate; [`PartialFiniteDifference`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/gradient_estimators.py) evaluates only selected coordinates and returns zero elsewhere.
+
+For a cheaper estimate in high dimension, [`SimultaneousPerturbation`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/gradient_estimators.py) draws $T$ Rademacher directions. At construction, the matrix $P\in\mathbb R^{T\times d}$ has entries $P_{ij}=s_{ij}/\sqrt T$, where each $s_{ij}$ is $+1$ or $-1$ with equal probability. The implementation uses
+
+```math
+\widehat{\nabla f}(\theta)
+=\frac{1}{\delta}P^\top
+\begin{pmatrix}
+f(\theta+\delta P_1)-f(\theta)\\
+\vdots\\
+f(\theta+\delta P_T)-f(\theta)
+\end{pmatrix}.
+```
+
+Because $\mathbb E[P^\top P]=I$, the first-order term recovers $\nabla f$ in expectation. Each update costs $T+1$ loss evaluations, independent of $d$ for a fixed $T$. Larger $T$ averages more directions but increases evaluation cost and memory.
+
+## 3. From an estimator to a model update
+
+The interfaces separate three operations:
+
+| Operation                                                  | Implementation                                                              |
+| ---------------------------------------------------------- | --------------------------------------------------------------------------- |
+| Generate perturbed parameters and reconstruct the gradient | [`gradient_estimators.py`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/gradient_estimators.py)     |
+| Generate Rademacher directions                             | [`perturbation_matrices.py`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/utils/perturbation_matrices.py)        |
+| Evaluate a model at nominal and perturbed parameters       | [`zeroth_order_blackbox.py`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/zeroth_order_blackbox.py) |
+| Apply SGD or Adam to the estimated gradient                | [`optimizers.py`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/optimizers.py)                       |
+
+The included [`ZerothOrderNeuralNetwork`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/neural_network/neural_network.py) is one black-box implementation. [`ParameterManager`](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/zeroth/zeroth_order/neural_network/parameter_manager.py) maps weights and biases to a flat $\theta$; NumPy broadcasting evaluates its $T+1$ perturbed networks in one batch. This batching reduces Python overhead, not the number of loss evaluations. Other models choose how to perform those evaluations.
+
+## 4. Numerical checks
+
+A deterministic [test](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/tests/test_gradient_estimator.py) applies the random-direction estimator to a linear function, whose gradient is known. The MNIST experiments provide end-to-end checks of zeroth-order training on CPU. Their protocols, results, plots, and raw loss traces are in [Experiments](https://github.com/nicolasmalet/Zeroth-Learn/blob/main/EXPERIMENTS.md).
+
+## 5. Reproduce the experiments
+
+With [uv](https://docs.astral.sh/uv/), the lockfile fixes the dependency versions. The first experiment run downloads MNIST from OpenML.
+
+```bash
+git clone https://github.com/nicolasmalet/Zeroth-Learn.git
+cd Zeroth-Learn
+uv sync --locked --extra mnist
+uv run --locked --extra mnist python -m unittest discover -s tests
+MPLBACKEND=Agg uv run --locked --extra mnist python -m lab.mnist
+MPLBACKEND=Agg uv run --locked --extra mnist python -m lab.mnist.run_experiment nb_perturbations_vs_model_size
+```
+
+The experiments write accuracy, loss traces, and plots under `results/`.
+
+Without uv, install the package with its MNIST extra in a virtual environment, then run the same Python modules.
+
+## 6. Scope
+
+The MNIST experiments validate classical models. They are not benchmarks of quantum hardware or of PyTorch, and their saved plots do not contain repeated trials or uncertainty estimates. Memory use grows with the number of perturbations in the vectorized neural-network implementation.

zeroth_learn-0.2.6/pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "zeroth-learn"
+version = "0.2.6"
+description = "Zeroth-order optimization for black-box models"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "numpy",
+    "pandas",
+    "matplotlib",
+    "cycler",
+]
+
+[project.optional-dependencies]
+mnist = ["scikit-learn"]
+
+[project.urls]
+Repository = "https://github.com/nicolasmalet/Zeroth-Learn"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["zeroth*"]
+exclude = ["lab*"]
+
+[tool.ruff]
+line-length = 120
+
+[tool.ruff.lint]
+select = ["E4", "E7", "E9", "F", "I", "UP", "B"]
+
+[tool.ruff.lint.per-file-ignores]
+"**/__init__.py" = ["F401", "F403"]
+"zeroth/all.py" = ["F403"]
+"lab/mnist/neural_networks.py" = ["E741"]

zeroth_learn-0.2.6/tests/test_gradient_estimator.py

@@ -0,0 +1,23 @@
+import unittest
+
+import numpy as np
+
+from zeroth.utils.perturbation_matrices import RademacherMatrix
+from zeroth.zeroth_order.gradient_estimators import SimultaneousPerturbation
+
+
+class SimultaneousPerturbationTest(unittest.TestCase):
+    def test_recovers_linear_gradient(self):
+        np.random.seed(42)
+        expected = np.array([1.5, -2.0, 0.5])
+        theta = np.array([0.2, -0.1, 0.3])
+        estimator = SimultaneousPerturbation(1e-6, 2_000, RademacherMatrix(), theta.size)
+
+        perturbed_theta = estimator.perturb(theta)
+        losses = (perturbed_theta @ expected)[:, None]
+
+        np.testing.assert_allclose(estimator.get_gradient(losses), expected, atol=0.06)
+
+
+if __name__ == "__main__":
+    unittest.main()

{zeroth_learn-0.2.4 → zeroth_learn-0.2.6}/zeroth/abstract/__init__.py

@@ -4,7 +4,7 @@ from .data_creator import DataCreator
 from .loss import Loss
 from .metric import Metric
 from .model import Model, ModelConfig
-from .neural_network import NeuralNetworkConfig, NeuralNetwork
+from .neural_network import NeuralNetwork, NeuralNetworkConfig
 from .optimizer import Optimizer
 from .perturbation_matrix import PerturbationMatrix
 from .summary import Summary

zeroth_learn-0.2.6/zeroth/abstract/activation.py

@@ -0,0 +1,20 @@
+from abc import ABC, abstractmethod
+
+from ..types import Array
+
+
+class Activation(ABC):
+    """Base class for activation functions."""
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+
+    @abstractmethod
+    def __call__(self, x: Array) -> Array:
+        """Apply the activation function."""
+        ...
+
+    @abstractmethod
+    def derivative(self, x: Array) -> Array:
+        """Compute the activation derivative."""
+        ...

{zeroth_learn-0.2.4 → zeroth_learn-0.2.6}/zeroth/abstract/blackbox.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from typing import Any
 
 from ..types import Array
 
@@ -24,7 +25,7 @@ class BlackBox(ABC):
         ...
 
     @abstractmethod
-    def init_params(self, params: dict) -> None:
+    def init_params(self, params: dict[str, Any]) -> None:
         """Manually initializes the weights and biases of the network.
 
         Args:
@@ -33,7 +34,7 @@ class BlackBox(ABC):
         ...
 
     @abstractmethod
-    def get_params(self) -> dict:
+    def get_params(self) -> dict[str, Any]:
         """Retrieves the current parameters of the network.
 
         Returns:

{zeroth_learn-0.2.4 → zeroth_learn-0.2.6}/zeroth/abstract/metric.py

@@ -9,5 +9,5 @@ class Metric(ABC):
 
     @abstractmethod
     def __call__(self, Y_pred: Array, Y_true: Array) -> float:
-        """Applique la fonction d'activation (forward)."""
+        """Compute the metric from predictions and targets."""
         ...

{zeroth_learn-0.2.4 → zeroth_learn-0.2.6}/zeroth/abstract/model.py

@@ -4,19 +4,21 @@ import pickle
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Callable
+from typing import Any
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+from matplotlib.figure import Figure
 
+from ..data import Data
+from ..plot_losses import plot_losses
+from ..types import Array
 from .blackbox import BlackBox
 from .loss import Loss
+from .metric import Metric
 from .optimizer import Optimizer
 from .summary import Summary
-from ..data import Data
-from ..plot_losses import plot_losses
-from ..types import Array
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -30,9 +32,9 @@ class ModelConfig(ABC, Summary):
     nb_epochs (int): Number of passes through the entire dataset.
     """
     name: str
-    id: dict
+    id: dict[str, Any]
     loss: Loss
-    metric: Callable
+    metric: Metric
     batch_size: int
     nb_epochs: int = 1
 
@@ -60,10 +62,10 @@ class Model(ABC):
         self.config = config
 
         self.name: str = config.name
-        self.id: dict = config.id
+        self.id: dict[str, Any] = config.id
         self.data: Data = data
         self.loss: Loss = config.loss
-        self.metric: Callable = config.metric
+        self.metric: Metric = config.metric
         self.batch_size: int = config.batch_size
         self.nb_epochs: int = config.nb_epochs
 
@@ -91,7 +93,7 @@ class Model(ABC):
         print_indexes = np.linspace(0, nb_batches - 1, nb_print).astype(int)
 
         for epoch_idx in range(self.nb_epochs):
-            print(f"        epoch n°{epoch_idx + 1} out of {self.nb_epochs}")
+            print(f"        epoch {epoch_idx + 1} of {self.nb_epochs}")
             self.data.permutation()
             self.data.batch_size = self.batch_size
 
@@ -100,12 +102,12 @@ class Model(ABC):
                 self.training_loss[epoch_idx * nb_batches + batch_idx] = avg_loss
 
                 if batch_idx in print_indexes:
-                    print(f"            batch n°{batch_idx + 1} out of {nb_batches}, "
+                    print(f"            batch {batch_idx + 1} of {nb_batches}, "
                           f"loss : {np.round(self.training_loss[epoch_idx * nb_batches + batch_idx], 3)}")
 
             self.test()
 
-    def plot_loss(self, smooth_fraction: float = 0.05) -> plt.Figure:
+    def plot_loss(self, smooth_fraction: float = 0.05) -> Figure:
         fig = plot_losses(dimension=0,
                           models=[self],
                           title=self.name,
@@ -113,7 +115,7 @@ class Model(ABC):
         plt.close(fig)
         return fig
 
-    def test(self) -> None:
+    def test(self) -> Array:
         X_test, Y_true = self.data.X_test, self.data.Y_test  # (in, batch), (out, batch)
         Y_pred = self.neural_network(X_test)  # (out, batch)
 
@@ -121,6 +123,7 @@ class Model(ABC):
         self.test_loss = self.loss.compute_loss(Y_pred, Y_true)
 
         print(f"    {self.id} accuracy : {self.test_accuracy}, loss : {self.test_loss}")
+        return Y_pred
 
     def save_loss(self, save_dir: Path) -> None:
         save_dir.mkdir(parents=True, exist_ok=True)
@@ -140,7 +143,7 @@ class Model(ABC):
             pickle.dump(params_dict, f)
 
     def load_weights(self, load_dir: Path) -> None:
-        """Restaure les paramètres depuis un fichier pickle."""
+        """Restore model parameters from a pickle file."""
         load_path = load_dir / self.WEIGHTS_FILE
         with load_path.open('rb') as f:
             params_dict = pickle.load(f)

{zeroth_learn-0.2.4 → zeroth_learn-0.2.6}/zeroth/abstract/neural_network.py

@@ -1,7 +1,7 @@
-from abc import ABC
 from dataclasses import dataclass
 
 from .activation import Activation
+from .blackbox import BlackBox
 from .summary import Summary
 
 
@@ -12,7 +12,7 @@ class NeuralNetworkConfig(Summary):
     activations: list[Activation]
 
 
-class NeuralNetwork(ABC):
+class NeuralNetwork(BlackBox):
     def __init__(self, config: NeuralNetworkConfig, input_dim: int, output_dim: int):
         self.name: str = config.name
         self.nb_layers: int = len(config.activations)

{zeroth_learn-0.2.4 → zeroth_learn-0.2.6}/zeroth/abstract/perturbation_matrix.py

@@ -4,12 +4,12 @@ from ..types import Array
 
 
 class PerturbationMatrix(ABC):
-    """Classe de base pour toutes les fonctions d'activation."""
+    """Base class for perturbation-matrix generators."""
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}()"
 
     @abstractmethod
     def __call__(self, nb_perturbation: int, nb_parameters: int) -> Array:
-        """Applique la fonction d'activation (forward)."""
+        """Generate perturbation directions."""
         ...

{zeroth_learn-0.2.4 → zeroth_learn-0.2.6}/zeroth/abstract/summary.py

@@ -16,17 +16,20 @@ class Summary:
             print(self._summary(self, indent=0), file=f)
 
     @classmethod
-    def load(cls, path: str, context: dict = None) -> Summary:
-        with open(path, "r") as f:
-            code = "".join(l for l in f)
-        return eval(code.strip(), {"__builtins__": __builtins__}, context)
+    def load(cls, path: str, context: dict[str, Any] | None = None) -> Summary:
+        with open(path) as f:
+            code = "".join(line for line in f)
+        result = eval(code.strip(), {"__builtins__": __builtins__}, context)
+        if not isinstance(result, cls):
+            raise TypeError(f"Expected {cls.__name__}, got {type(result).__name__}")
+        return result
 
     def _summary(self, obj: Any, indent: int) -> str:
         shift = "    " * indent
         next_shift = "    " * (indent + 1)
 
         if dataclasses.is_dataclass(obj):
-            cls_name = obj.__class__.__name__
+            cls_name = type(obj).__name__
             items = []
             for f in dataclasses.fields(obj):
                 val = getattr(obj, f.name)
@@ -37,13 +40,15 @@ class Summary:
             return f"{cls_name}(\n{content}\n{shift})"
 
         elif isinstance(obj, list):
-            if not obj: return "[]"
+            if not obj:
+                return "[]"
             items = [self._summary(item, indent + 1) for item in obj]
             content = ",\n".join([f"{next_shift}{item.lstrip()}" for item in items])
             return f"[\n{content}\n{shift}]"
 
         elif isinstance(obj, dict):
-            if not obj: return "{}"
+            if not obj:
+                return "{}"
             items = [f"{next_shift}{repr(k)}: {self._summary(v, indent + 1).lstrip()}" for k, v in obj.items()]
             content = ",\n".join(items)
             return f"{{\n{content}\n{shift}}}"

zeroth_learn-0.2.6/zeroth/all.py

@@ -0,0 +1,5 @@
+from .experiment import *
+from .first_order import *
+from .losses import *
+from .utils import *
+from .zeroth_order import *

{zeroth_learn-0.2.4 → zeroth_learn-0.2.6}/zeroth/data.py

@@ -1,4 +1,4 @@
-from typing import Iterator
+from collections.abc import Iterator
 
 import numpy as np
 
@@ -17,7 +17,7 @@ class Data:
         self.input_dim: int = self.raw_X_train.shape[1]
         self.output_dim: int = self.raw_Y_train.shape[1] if nb_class == 0 else nb_class
 
-        self.batch_size: int | None = None
+        self.batch_size: int = 1
 
         self.indices: np.ndarray = np.arange(self.nb_data)
 

{zeroth_learn-0.2.4 → zeroth_learn-0.2.6}/zeroth/experiment.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
 import itertools
-from pathlib import Path
 from dataclasses import dataclass, replace
-from typing import Union
+from pathlib import Path
+from typing import Any, cast
 
 import matplotlib.pyplot as plt
 import pandas as pd
+from matplotlib.figure import Figure
 
-from .abstract import Model, ModelConfig, DataCreator, Summary
+from .abstract import DataCreator, Model, ModelConfig, Summary
 from .data import Data
 from .plot_losses import plot_losses
 from .utils.dataclasses_utils import get_name, set_value_by_path
@@ -18,7 +19,7 @@ from .utils.dataclasses_utils import get_name, set_value_by_path
 class VariationConfig:
     name: str
     param: list[str]
-    values: Union[list, list[list]]
+    values: list[list[Any]]
 
 
 @dataclass(frozen=True)
@@ -54,11 +55,11 @@ class Experiment:
         self.models: list[Model] = generate_models(config.base_model, config.variations, self.data)
 
     def train_models(self, nb_print: int) -> None:
-        print(f"Training Models")
+        print("Training Models")
         for model in self.models:
             model.train(nb_print)
 
-    def plot_losses(self, title: str, plot_dimension: int, smooth_fraction: float = 0) -> plt.Figure:
+    def plot_losses(self, title: str, plot_dimension: int, smooth_fraction: float = 0) -> Figure:
         fig = plot_losses(title=title,
                           dimension=plot_dimension,
                           models=self.models,
@@ -69,7 +70,7 @@ class Experiment:
         return fig
 
     def test_models(self) -> None:
-        print(f"Testing Models")
+        print("Testing Models")
         for model in self.models:
             model.test()
 
@@ -88,7 +89,7 @@ class Experiment:
         df.to_csv(save_path)
 
     def save_weights(self, save_dir: Path) -> None:
-        for i, model in enumerate(self.models):
+        for model in self.models:
             save_path = save_dir / model.name
             model.save_weights(save_path)
 
@@ -97,7 +98,7 @@ class Experiment:
         config_path = save_dir / self.CONFIG_FILE
         self.config.save(config_path)
 
-        for i, model in enumerate(self.models):
+        for model in self.models:
             save_path = save_dir / model.name
             model.config.save(save_path)
 
@@ -109,13 +110,13 @@ def generate_models(base_model: ModelConfig, variations: list[VariationConfig],
 
     for combination in itertools.product(*values_lists):
         id_ = {}
-        current_model = base_model
+        current_model: ModelConfig = base_model
 
-        for var_config, current_vals in zip(variations, combination):
+        for var_config, current_vals in zip(variations, combination, strict=True):
             id_[var_config.name] = get_name(current_vals[0])
 
-            for path, val in zip(var_config.param, current_vals):
-                current_model = set_value_by_path(current_model, path, val)
+            for path, val in zip(var_config.param, current_vals, strict=True):
+                current_model = cast(ModelConfig, set_value_by_path(current_model, path, val))
 
         current_model = replace(current_model, id=id_)
         models.append(current_model.instantiate(data))

zeroth_learn-0.2.6/zeroth/first_order/__init__.py

@@ -0,0 +1,4 @@
+from .layer import Layer
+from .model import FirstOrderModel, FirstOrderModelConfig
+from .neural_network import FirstOrderNeuralNetwork
+from .optimizers import FirstOrderAdamConfig, FirstOrderOptimizer, FirstOrderOptimizerConfig, FirstOrderSGDConfig

{zeroth_learn-0.2.4 → zeroth_learn-0.2.6}/zeroth/first_order/model.py

@@ -2,10 +2,10 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
-from .neural_network import FirstOrderNeuralNetwork
-from .optimizers import FirstOrderOptimizerConfig, FirstOrderOptimizer
-from ..abstract import ModelConfig, Model, NeuralNetworkConfig
+from ..abstract import Model, ModelConfig, NeuralNetworkConfig
 from ..data import Data
+from .neural_network import FirstOrderNeuralNetwork
+from .optimizers import FirstOrderOptimizerConfig
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -21,6 +21,5 @@ class FirstOrderModel(Model):
     def __init__(self, config: FirstOrderModelConfig, data: Data) -> None:
         super().__init__(config, data)
 
-        self.neural_network: FirstOrderNeuralNetwork = FirstOrderNeuralNetwork(config.neural_network_config,
-                                                                               data.input_dim, data.output_dim)
-        self.optimizer: FirstOrderOptimizer = config.optimizer_config.instantiate()
+        self.neural_network = FirstOrderNeuralNetwork(config.neural_network_config, data.input_dim, data.output_dim)
+        self.optimizer = config.optimizer_config.instantiate()

{zeroth_learn-0.2.4 → zeroth_learn-0.2.6}/zeroth/first_order/neural_network.py

@@ -1,6 +1,8 @@
-from .layer import Layer
+from typing import Any
+
 from ..abstract.neural_network import NeuralNetwork, NeuralNetworkConfig
 from ..types import Array
+from .layer import Layer
 
 
 class FirstOrderNeuralNetwork(NeuralNetwork):
@@ -25,14 +27,14 @@ class FirstOrderNeuralNetwork(NeuralNetwork):
     def __call__(self, X: Array) -> Array:
         return self.forward(X)
 
-    def init_params(self, params: dict) -> None:
+    def init_params(self, params: dict[str, Any]) -> None:
         Ws = params["Ws"]
         Bs = params["Bs"]
-        for layer, W, B in zip(self.layers, Ws, Bs):
+        for layer, W, B in zip(self.layers, Ws, Bs, strict=True):
             layer.W = W
             layer.B = B
 
-    def get_params(self) -> dict:
+    def get_params(self) -> dict[str, Any]:
         Ws, Bs = [], []
         for layer in self.layers:
             Ws.append(layer.W)