From c593091d60f65a7e8dde48a56734fcf60968bb0e Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 13:02:10 -0400 Subject: [PATCH 01/41] Remove linting from main ci workflow --- .github/workflows/run_tests.yml | 28 ++++++---------------------- 1 file changed, 6 insertions(+), 22 deletions(-) diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index 23e1d2a..d97d32a 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -19,24 +19,19 @@ jobs: with: persist-credentials: false - - name: Setup Python - uses: actions/setup-python@v6 - with: - python-version: ${{ matrix.python-version }} - - name: Install uv uses: astral-sh/setup-uv@v7 with: - version: "0.6.5" + version: "0.12.0" # enable-cache: true # cache-dependency-glob: "pyproject.toml pdm.lock" - name: Clear all caches run: | - rm -rf ~/.cache/pip - rm -rf ~/.cache/uv - rm -rf ~/.cache/conda - rm -rf ~/.cache/npm + rm -rf ~/.cache/pip + rm -rf ~/.cache/uv + rm -rf ~/.cache/conda + rm -rf ~/.cache/npm - name: Install package run: uv sync --all-groups --reinstall @@ -47,12 +42,6 @@ jobs: - name: Run pytest run: uv run pytest - - name: Check styling - run: uv run ruff format --check . - - - name: Linting - run: uv run ruff check src/lanfactory - - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v7 with: @@ -67,15 +56,10 @@ jobs: with: persist-credentials: false - - name: Setup Python - uses: actions/setup-python@v6 - with: - python-version: "3.12" - - name: Install uv uses: astral-sh/setup-uv@v7 with: - version: "0.6.5" + version: "0.12.0" - name: Install package (with notebook + backend deps) run: uv sync --all-groups From b4e44c16c0098ec308395ab1a86a816fd0897e78 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 13:02:31 -0400 Subject: [PATCH 02/41] Add linting and formatting workflow --- .github/workflows/linting_formatting.yml | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 .github/workflows/linting_formatting.yml diff --git a/.github/workflows/linting_formatting.yml b/.github/workflows/linting_formatting.yml new file mode 100644 index 0000000..fffcca3 --- /dev/null +++ b/.github/workflows/linting_formatting.yml @@ -0,0 +1,24 @@ +name: Linting and Formatting + +on: + pull_request: + +jobs: + lint: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v7 + with: + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v7 + with: + version: "0.12.0" + + - name: Check styling + run: uvx ruff format --check . + + - name: Check linting + run: uvx ruff check src/lanfactory \ No newline at end of file From b72172c3f9707d69123746f5561e18f95ab717e1 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 13:04:54 -0400 Subject: [PATCH 03/41] Fix formatting --- .github/workflows/linting_formatting.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/linting_formatting.yml b/.github/workflows/linting_formatting.yml index fffcca3..bed620a 100644 --- a/.github/workflows/linting_formatting.yml +++ b/.github/workflows/linting_formatting.yml @@ -18,7 +18,7 @@ jobs: version: "0.12.0" - name: Check styling - run: uvx ruff format --check . + run: uvx ruff format --check . - - name: Check linting - run: uvx ruff check src/lanfactory \ No newline at end of file + - name: Check linting + run: uvx ruff check src/lanfactory From e9ffbbdbbfe16c8a366ea575f11812d379f7d69c Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 13:10:11 -0400 Subject: [PATCH 04/41] Autofix fixable linting problems --- src/lanfactory/__init__.py | 7 ++---- src/lanfactory/cli/download_hf.py | 1 - src/lanfactory/cli/jax_train.py | 6 +++-- src/lanfactory/cli/torch_train.py | 9 +++---- src/lanfactory/cli/upload_hf.py | 1 - src/lanfactory/cli/utils.py | 7 +++--- src/lanfactory/config/__init__.py | 16 ++++++------- src/lanfactory/hf/__init__.py | 16 ++++++------- src/lanfactory/hf/upload.py | 3 ++- src/lanfactory/network_inspectors/__init__.py | 6 ++--- src/lanfactory/network_inspectors/api.py | 2 +- src/lanfactory/network_inspectors/config.py | 1 - src/lanfactory/network_inspectors/loaders.py | 4 ++-- src/lanfactory/onnx/__init__.py | 4 ++-- src/lanfactory/trainers/__init__.py | 8 +++---- src/lanfactory/trainers/jax_mlp.py | 9 +++---- src/lanfactory/trainers/torch_mlp.py | 24 ++++++++----------- src/lanfactory/utils/__init__.py | 4 ++-- src/lanfactory/utils/mlflow_utils.py | 3 ++- 19 files changed, 64 insertions(+), 67 deletions(-) diff --git a/src/lanfactory/__init__.py b/src/lanfactory/__init__.py index 780a2f4..33b3b8f 100755 --- a/src/lanfactory/__init__.py +++ b/src/lanfactory/__init__.py @@ -1,11 +1,8 @@ __version__ = "0.8.0" -from . import config -from . import trainers -from . import utils -from . import onnx +from . import config, onnx, trainers, utils -__all__ = ["config", "trainers", "utils", "onnx", "network_inspectors"] +__all__ = ["config", "network_inspectors", "onnx", "trainers", "utils"] def __getattr__(name): diff --git a/src/lanfactory/cli/download_hf.py b/src/lanfactory/cli/download_hf.py index 262f776..7001b45 100644 --- a/src/lanfactory/cli/download_hf.py +++ b/src/lanfactory/cli/download_hf.py @@ -12,7 +12,6 @@ from pathlib import Path import typer - from lanfactory.hf import DEFAULT_REPO_ID, VALID_NETWORK_TYPES app = typer.Typer() diff --git a/src/lanfactory/cli/jax_train.py b/src/lanfactory/cli/jax_train.py index ffde030..dedcc58 100644 --- a/src/lanfactory/cli/jax_train.py +++ b/src/lanfactory/cli/jax_train.py @@ -164,9 +164,10 @@ def main( if mlflow_tracking_enabled: try: - import mlflow import os + import mlflow + # Set tracking URI with priority: CLI arg > env var > default if mlflow_tracking_uri: tracking_uri = mlflow_tracking_uri @@ -234,9 +235,10 @@ def main( if not mlflow_tracking_enabled: # Need to initialize MLflow just for querying - import mlflow import os + import mlflow + # Use same logic as above for tracking URI if mlflow_tracking_uri: tracking_uri = mlflow_tracking_uri diff --git a/src/lanfactory/cli/torch_train.py b/src/lanfactory/cli/torch_train.py index 84b1855..33eb128 100644 --- a/src/lanfactory/cli/torch_train.py +++ b/src/lanfactory/cli/torch_train.py @@ -22,11 +22,10 @@ from importlib.resources import as_file, files from pathlib import Path +import lanfactory import psutil import torch import typer - -import lanfactory from lanfactory.cli.utils import ( _get_train_network_config, ) @@ -168,9 +167,10 @@ def main( if mlflow_tracking_enabled: try: - import mlflow import os + import mlflow + # Set tracking URI with priority: CLI arg > env var > default if mlflow_tracking_uri: tracking_uri = mlflow_tracking_uri @@ -238,9 +238,10 @@ def main( if not mlflow_tracking_enabled: # Need to initialize MLflow just for querying - import mlflow import os + import mlflow + # Use same logic as above for tracking URI if mlflow_tracking_uri: tracking_uri = mlflow_tracking_uri diff --git a/src/lanfactory/cli/upload_hf.py b/src/lanfactory/cli/upload_hf.py index 538b1db..4142a83 100644 --- a/src/lanfactory/cli/upload_hf.py +++ b/src/lanfactory/cli/upload_hf.py @@ -12,7 +12,6 @@ from pathlib import Path import typer - from lanfactory.hf import DEFAULT_REPO_ID, VALID_NETWORK_TYPES app = typer.Typer() diff --git a/src/lanfactory/cli/utils.py b/src/lanfactory/cli/utils.py index 7933973..f3bad93 100644 --- a/src/lanfactory/cli/utils.py +++ b/src/lanfactory/cli/utils.py @@ -1,10 +1,11 @@ # import argparse import logging -from pathlib import Path import pickle -import yaml -import numpy as np +from pathlib import Path + import lanfactory +import numpy as np +import yaml logger = logging.getLogger(__name__) diff --git a/src/lanfactory/config/__init__.py b/src/lanfactory/config/__init__.py index 1878766..42aee52 100755 --- a/src/lanfactory/config/__init__.py +++ b/src/lanfactory/config/__init__.py @@ -1,21 +1,21 @@ from .network_configs import ( - network_config_mlp, network_config_choice_prob, - network_config_opn, network_config_cpn, - train_config_mlp, + network_config_mlp, + network_config_opn, train_config_choice_prob, - train_config_opn, train_config_cpn, + train_config_mlp, + train_config_opn, ) __all__ = [ - "network_config_mlp", "network_config_choice_prob", - "network_config_opn", "network_config_cpn", - "train_config_mlp", + "network_config_mlp", + "network_config_opn", "train_config_choice_prob", - "train_config_opn", "train_config_cpn", + "train_config_mlp", + "train_config_opn", ] diff --git a/src/lanfactory/hf/__init__.py b/src/lanfactory/hf/__init__.py index 7dbda0a..f83e6c5 100644 --- a/src/lanfactory/hf/__init__.py +++ b/src/lanfactory/hf/__init__.py @@ -7,20 +7,20 @@ DEFAULT_REPO_ID = "franklab/HSSM" VALID_NETWORK_TYPES = ("lan", "cpn", "opn") -from lanfactory.hf.model_card import ( # noqa: E402 - load_model_card_yaml, - generate_readme, +from lanfactory.hf.download import download_model +from lanfactory.hf.model_card import ( ModelCardConfig, + generate_readme, + load_model_card_yaml, ) -from lanfactory.hf.upload import upload_model # noqa: E402 -from lanfactory.hf.download import download_model # noqa: E402 +from lanfactory.hf.upload import upload_model __all__ = [ "DEFAULT_REPO_ID", "VALID_NETWORK_TYPES", - "load_model_card_yaml", - "generate_readme", "ModelCardConfig", - "upload_model", "download_model", + "generate_readme", + "load_model_card_yaml", + "upload_model", ] diff --git a/src/lanfactory/hf/upload.py b/src/lanfactory/hf/upload.py index c223ef2..adbf23f 100644 --- a/src/lanfactory/hf/upload.py +++ b/src/lanfactory/hf/upload.py @@ -155,7 +155,8 @@ def _upload_to_hf( # pragma: no cover ) -> str: """HF-dependent implementation of upload_model.""" try: - from huggingface_hub import HfApi, create_repo as hf_create_repo + from huggingface_hub import HfApi + from huggingface_hub import create_repo as hf_create_repo except ImportError as exc: raise ImportError( "huggingface_hub is required for HuggingFace uploads. " diff --git a/src/lanfactory/network_inspectors/__init__.py b/src/lanfactory/network_inspectors/__init__.py index 6f8e2ab..6154967 100644 --- a/src/lanfactory/network_inspectors/__init__.py +++ b/src/lanfactory/network_inspectors/__init__.py @@ -15,10 +15,10 @@ from .loaders import get_torch_mlp __all__ = [ + "GridSpec", + "ModelSpec", + "PlotConfig", "get_torch_mlp", "kde_vs_lan_likelihoods", "lan_manifold", - "ModelSpec", - "PlotConfig", - "GridSpec", ] diff --git a/src/lanfactory/network_inspectors/api.py b/src/lanfactory/network_inspectors/api.py index 0635786..570cdbe 100644 --- a/src/lanfactory/network_inspectors/api.py +++ b/src/lanfactory/network_inspectors/api.py @@ -2,8 +2,8 @@ from __future__ import annotations -from collections.abc import Callable import logging +from collections.abc import Callable from typing import TYPE_CHECKING, Any import numpy as np diff --git a/src/lanfactory/network_inspectors/config.py b/src/lanfactory/network_inspectors/config.py index 1eee71c..2c621e6 100644 --- a/src/lanfactory/network_inspectors/config.py +++ b/src/lanfactory/network_inspectors/config.py @@ -8,7 +8,6 @@ import numpy as np from numpy.typing import NDArray - from ssms.config import ModelConfigBuilder diff --git a/src/lanfactory/network_inspectors/loaders.py b/src/lanfactory/network_inspectors/loaders.py index 4940051..88fa3a5 100644 --- a/src/lanfactory/network_inspectors/loaders.py +++ b/src/lanfactory/network_inspectors/loaders.py @@ -2,11 +2,11 @@ from __future__ import annotations -from collections.abc import Callable import os +import warnings +from collections.abc import Callable from os import PathLike from typing import Any -import warnings import numpy as np from numpy.typing import NDArray diff --git a/src/lanfactory/onnx/__init__.py b/src/lanfactory/onnx/__init__.py index 0b25cb2..289ac30 100755 --- a/src/lanfactory/onnx/__init__.py +++ b/src/lanfactory/onnx/__init__.py @@ -3,7 +3,7 @@ from .transform_onnx import transform_to_onnx __all__ = [ - "transform_to_onnx", - "transform_sbi_to_onnx", "transform_bayesflow_to_onnx", + "transform_sbi_to_onnx", + "transform_to_onnx", ] diff --git a/src/lanfactory/trainers/__init__.py b/src/lanfactory/trainers/__init__.py index ecd495d..6ba03f6 100755 --- a/src/lanfactory/trainers/__init__.py +++ b/src/lanfactory/trainers/__init__.py @@ -1,16 +1,16 @@ import warnings +from .jax_mlp import JaxMLP, JaxMLPFactory, ModelTrainerJaxMLP from .torch_mlp import ( DatasetTorch, - TorchMLP, - TorchMLPFactory, - ModelTrainerTorchMLP, LoadTorchMLP, LoadTorchMLPInfer, + ModelTrainerTorchMLP, + TorchMLP, + TorchMLPFactory, make_dataloader, make_train_valid_dataloaders, ) -from .jax_mlp import JaxMLPFactory, JaxMLP, ModelTrainerJaxMLP __all__ = [ # Dataset and DataLoader helpers diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 3c3968c..430d7c8 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -3,10 +3,11 @@ """ import pickle +from collections.abc import Callable, Sequence from functools import partial from pathlib import Path from time import time -from typing import Any, Callable, Sequence +from typing import Any import flax import jax @@ -134,7 +135,7 @@ def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray: if (not self.train) and (self.train_output_type == "logprob"): x = x # just for pedagogy elif (not self.train) and (self.train_output_type == "logits"): - x = -jnp.log((1 + jnp.exp(-x))) + x = -jnp.log(1 + jnp.exp(-x)) elif not self.train: # pragma: no cover x = x # just for pedagogy @@ -272,7 +273,7 @@ def __init__( The ModelTrainerJaxMLP object. """ - if "loss_dict" not in train_config.keys(): + if "loss_dict" not in train_config: self.loss_dict: dict[str, dict] = { "huber": {"fun": optax.huber_loss, "kwargs": {"delta": 1}}, "mse": {"fun": optax.l2_loss, "kwargs": {}}, @@ -281,7 +282,7 @@ def __init__( else: # pragma: no cover self.loss_dict = train_config["loss_dict"] - if "lr_dict" not in train_config.keys(): + if "lr_dict" not in train_config: # Todo: Add more schedules (for now warmup_cosine_decay_schedule) self.lr_dict: dict[str, float] = { "init_value": 0.0002, diff --git a/src/lanfactory/trainers/torch_mlp.py b/src/lanfactory/trainers/torch_mlp.py index d439bb7..8dc34be 100755 --- a/src/lanfactory/trainers/torch_mlp.py +++ b/src/lanfactory/trainers/torch_mlp.py @@ -1,18 +1,16 @@ """This module contains the classes for training TorchMLP models.""" -import numpy as np -import pandas as pd -import pickle -from typing import Callable -from time import time import logging +import pickle +from collections.abc import Callable from pathlib import Path +from time import time - +import numpy as np +import pandas as pd import torch -import torch.nn as nn -import torch.optim as optim import torch.nn.functional as F +from torch import nn, optim from torch.utils.data import DataLoader try: @@ -100,7 +98,6 @@ def __load_file(self, file_index: int) -> None: shuffle_idx, : ] self.tmp_data[self.label_key] = self.tmp_data[self.label_key][shuffle_idx] - return def __init_file_shape(self) -> None: # Function gets dimensionalities form a test data file @@ -130,7 +127,6 @@ def __init_file_shape(self) -> None: self.label_dim = self.file_shape_dict["labels"][1] else: self.label_dim = 1 - return def __data_generation( self, batch_ids: np.ndarray | None = None @@ -337,7 +333,7 @@ def TorchMLPFactory( """ if isinstance(network_config, str): with open(network_config, "rb") as f: - network_config = pickle.load(f) # noqa: S301 + network_config = pickle.load(f) assert isinstance(network_config, dict) return TorchMLP( @@ -366,7 +362,7 @@ def __init__( input_shape: int = 10, network_type: str | None = None, ) -> None: - super(TorchMLP, self).__init__() + super().__init__() self.input_shape = input_shape self.network_config = network_config @@ -452,7 +448,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.layers[-1](x) elif self.train_output_type == "logits": return -torch.log( - (1 + torch.exp(-self.layers[-1](x))) + 1 + torch.exp(-self.layers[-1](x)) ) # log ( 1 / (1 + exp(-x))), where x = log(p / (1 - p)) else: return self.layers[-1](x) @@ -501,7 +497,7 @@ def __init__( self.train_config: dict = pickle.load(open(train_config, "rb")) except (OSError, pickle.PickleError) as e: # pragma: no cover logger.error( - f"Error loading training config from file {train_config}: {str(e)}" + f"Error loading training config from file {train_config}: {e!s}" ) raise elif isinstance(train_config, dict): diff --git a/src/lanfactory/utils/__init__.py b/src/lanfactory/utils/__init__.py index 0e9effc..029cb29 100755 --- a/src/lanfactory/utils/__init__.py +++ b/src/lanfactory/utils/__init__.py @@ -1,11 +1,11 @@ -from .util_funs import save_configs from .mlflow_utils import ( get_files_from_data_generation_experiment, log_training_data_lineage, ) +from .util_funs import save_configs __all__ = [ - "save_configs", "get_files_from_data_generation_experiment", "log_training_data_lineage", + "save_configs", ] diff --git a/src/lanfactory/utils/mlflow_utils.py b/src/lanfactory/utils/mlflow_utils.py index 2d93193..d02a470 100644 --- a/src/lanfactory/utils/mlflow_utils.py +++ b/src/lanfactory/utils/mlflow_utils.py @@ -141,8 +141,9 @@ def log_training_data_lineage( - extra_files """ try: - import mlflow import os + + import mlflow except ImportError: logger.error("mlflow package not installed") return {} From 08d232414454f79012d4167d5159b53a5c2c8233 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 13:10:22 -0400 Subject: [PATCH 05/41] Use context managers --- src/lanfactory/utils/util_funs.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/lanfactory/utils/util_funs.py b/src/lanfactory/utils/util_funs.py index 6718380..589a5df 100755 --- a/src/lanfactory/utils/util_funs.py +++ b/src/lanfactory/utils/util_funs.py @@ -28,14 +28,13 @@ def save_configs( Path(save_folder).mkdir(parents=True, exist_ok=True) # Save network config - pickle.dump( - network_config, - open(Path(save_folder) / f"{model_id}_network_config.pickle", "wb"), - ) + network_config_path = Path(save_folder) / f"{model_id}_network_config.pickle" + with open(network_config_path, "wb") as f: + pickle.dump(network_config, f) print("Saved network config") + # Save train config - pickle.dump( - train_config, open(Path(save_folder) / f"{model_id}_train_config.pickle", "wb") - ) + train_config_path = Path(save_folder) / f"{model_id}_train_config.pickle" + with open(train_config_path, "wb") as f: + pickle.dump(train_config, f) print("Saved train config") - return From f2d3613933cca1b8bb5a4443eaef44bdffa7a1b2 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 13:14:35 -0400 Subject: [PATCH 06/41] Rewrite as set comprehension --- src/lanfactory/utils/mlflow_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lanfactory/utils/mlflow_utils.py b/src/lanfactory/utils/mlflow_utils.py index d02a470..a327460 100644 --- a/src/lanfactory/utils/mlflow_utils.py +++ b/src/lanfactory/utils/mlflow_utils.py @@ -166,7 +166,7 @@ def log_training_data_lineage( # Verify we have all the files we expect expected_file_names = set(expected_files_info["all_files"]) - actual_file_names = set(f.name for f in valid_file_list) + actual_file_names = {f.name for f in valid_file_list} missing_files = expected_file_names - actual_file_names extra_files = actual_file_names - expected_file_names From 5f725f1531c23f45d51e4c0396a60f854b62a506 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 13:18:41 -0400 Subject: [PATCH 07/41] Update type in function signature arg to include None --- src/lanfactory/utils/mlflow_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lanfactory/utils/mlflow_utils.py b/src/lanfactory/utils/mlflow_utils.py index a327460..d9de6bf 100644 --- a/src/lanfactory/utils/mlflow_utils.py +++ b/src/lanfactory/utils/mlflow_utils.py @@ -110,7 +110,7 @@ def log_training_data_lineage( training_data_folder: Path, valid_file_list: list, n_training_files: int, - tracking_uri: str = None, + tracking_uri: str | None = None, ) -> dict: """Log training data lineage information to MLflow. From fcddb5a600f7aedc4d063095c6f1760428b9ffb0 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 13:18:57 -0400 Subject: [PATCH 08/41] Use enumerate --- src/lanfactory/trainers/torch_mlp.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/lanfactory/trainers/torch_mlp.py b/src/lanfactory/trainers/torch_mlp.py index 8dc34be..5e9779c 100755 --- a/src/lanfactory/trainers/torch_mlp.py +++ b/src/lanfactory/trainers/torch_mlp.py @@ -661,7 +661,7 @@ def train_and_evaluate( epoch_s_t = time() # Training loop - for xb, yb in self.train_dl: + for cnt, (xb, yb) in enumerate(self.train_dl): # Shift data to device if self.pin_memory and str(self.dev) == "cuda": xb, yb = xb.cuda(non_blocking=True), yb.cuda(non_blocking=True) @@ -680,7 +680,6 @@ def train_and_evaluate( # Log training progress self._log_training_progress(epoch, cnt, loss, verbose) - cnt += 1 step_cnt += 1 print( From 96072772af44c9ccc0310dd935f333ac0d97dff1 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 13:20:58 -0400 Subject: [PATCH 09/41] Remove unused variable --- src/lanfactory/trainers/jax_mlp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 430d7c8..04d5084 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -562,7 +562,7 @@ def train_and_evaluate( # Training loop over epochs for epoch in range(self.train_config["n_epochs"]): print("Epoch: " + str(epoch) + " of " + str(self.train_config["n_epochs"])) - state, train_loss = self.run_epoch( + state, _ = self.run_epoch( state, train=True, verbose=verbose, From faf7f0b5b5d238ea4a23693a1acba7aeb442f96e Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:31:51 -0400 Subject: [PATCH 10/41] Use get method --- src/lanfactory/trainers/torch_mlp.py | 22 +++++----------------- 1 file changed, 5 insertions(+), 17 deletions(-) diff --git a/src/lanfactory/trainers/torch_mlp.py b/src/lanfactory/trainers/torch_mlp.py index 5e9779c..e828086 100755 --- a/src/lanfactory/trainers/torch_mlp.py +++ b/src/lanfactory/trainers/torch_mlp.py @@ -574,36 +574,24 @@ def __get_scheduler(self) -> None: self.optimizer, mode="min", factor=( - self.train_config["lr_scheduler_params"]["factor"] - if "factor" in self.train_config["lr_scheduler_params"] - else 0.1 + self.train_config["lr_scheduler_params"].get("factor", 0.1) ), patience=( - self.train_config["lr_scheduler_params"]["patience"] - if "patience" in self.train_config["lr_scheduler_params"] - else 2 + self.train_config["lr_scheduler_params"].get("patience", 2) ), threshold=( - self.train_config["lr_scheduler_params"]["threshold"] - if "threshold" in self.train_config["lr_scheduler_params"] - else 0.001 + self.train_config["lr_scheduler_params"].get("threshold", 0.001) ), threshold_mode="rel", cooldown=0, min_lr=( - self.train_config["lr_scheduler_params"]["min_lr"] - if "min_lr" in self.train_config["lr_scheduler_params"] - else 0.00000001 + self.train_config["lr_scheduler_params"].get("min_lr", 1e-08) ), ) elif self.train_config["lr_scheduler"] == "multiply": self.scheduler = optim.lr_scheduler.ExponentialLR( self.optimizer, - gamma=( - self.train_config["lr_scheduler_params"]["factor"] - if "factor" in self.train_config["lr_scheduler_params"] - else 0.1 - ), + gamma=(self.train_config["lr_scheduler_params"].get("factor", 0.1)), last_epoch=-1, ) From e1f9c4f54b887f31979fea220053b410a9ae4c6b Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:41:39 -0400 Subject: [PATCH 11/41] Update torch_mlp --- src/lanfactory/trainers/torch_mlp.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/lanfactory/trainers/torch_mlp.py b/src/lanfactory/trainers/torch_mlp.py index e828086..8fc9708 100755 --- a/src/lanfactory/trainers/torch_mlp.py +++ b/src/lanfactory/trainers/torch_mlp.py @@ -88,7 +88,8 @@ def __getitem__(self, index: int) -> tuple[np.ndarray, np.ndarray]: def __load_file(self, file_index: int) -> None: # Load file and shuffle the indices - self.tmp_data = pickle.load(open(self.file_ids[file_index], "rb")) + with open(self.file_ids[file_index], "rb") as f: + self.tmp_data = pickle.load(f) shuffle_idx = np.random.choice( self.tmp_data[self.features_key].shape[0], size=self.tmp_data[self.features_key].shape[0], @@ -101,7 +102,8 @@ def __load_file(self, file_index: int) -> None: def __init_file_shape(self) -> None: # Function gets dimensionalities form a test data file - init_file = pickle.load(open(self.file_ids[0], "rb")) + with open(self.file_ids[0], "rb") as f: + init_file = pickle.load(f) self.file_shape_dict = { "inputs": init_file[self.features_key].shape, "labels": init_file[self.label_key].shape, @@ -493,11 +495,15 @@ def __init__( elif isinstance(train_config, str | Path): print("train_config is passed as string or path: \n", train_config) try: - logger.info("Trying to load string as path to pickle file: ") - self.train_config: dict = pickle.load(open(train_config, "rb")) + logger.info( + "Trying to load string as path to pickle file: %s", train_config + ) + self.train_config: dict = pickle.loads(Path(train_config).read_bytes()) except (OSError, pickle.PickleError) as e: # pragma: no cover logger.error( - f"Error loading training config from file {train_config}: {e!s}" + "Error loading training config from file %s: %s", + train_config, + e, ) raise elif isinstance(train_config, dict): From 283e2fc45b78daaf4d5d0d0681fab6a493ae67a9 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:41:54 -0400 Subject: [PATCH 12/41] Use set comp --- src/lanfactory/cli/jax_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lanfactory/cli/jax_train.py b/src/lanfactory/cli/jax_train.py index dedcc58..faf15fd 100644 --- a/src/lanfactory/cli/jax_train.py +++ b/src/lanfactory/cli/jax_train.py @@ -314,7 +314,7 @@ def main( # Mode 2: Validation - verify MLflow files exist in training_data_folder if mlflow_lineage_info and training_data_folder: expected_files = set(mlflow_lineage_info["all_files"]) - actual_files = set(f.name for f in valid_file_list) + actual_files = {f.name for f in valid_file_list} missing_files = expected_files - actual_files extra_files = actual_files - expected_files From 05dadd208ec6dfa8e7b7f5a6fabc77ee02397c91 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:42:10 -0400 Subject: [PATCH 13/41] Simplify serialization step --- src/lanfactory/cli/jax_train.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/lanfactory/cli/jax_train.py b/src/lanfactory/cli/jax_train.py index faf15fd..0132ef7 100644 --- a/src/lanfactory/cli/jax_train.py +++ b/src/lanfactory/cli/jax_train.py @@ -465,13 +465,8 @@ def main( ] ) - pickle.dump( - network_config, - open( - networks_path / file_name_suffix, - "wb", - ), - ) + file_path = networks_path / file_name_suffix + file_path.write_bytes(pickle.dumps(network_config)) # Load network net = lanfactory.trainers.JaxMLPFactory( From 53d9e9de2ef34cfa9207317862bdbd7ee2a95f56 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:42:33 -0400 Subject: [PATCH 14/41] Delete commeted code --- src/lanfactory/trainers/jax_mlp.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 04d5084..20c48ce 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -86,10 +86,6 @@ class JaxMLP(nn.Module): activations_dict = frozendict( {"relu": nn.relu, "tanh": nn.tanh, "sigmoid": nn.sigmoid} ) - # network_type: Optional[str] = "none" - - # Define network type - # network_type = "lan" if train_output_type == "logprob" else "cpn" def setup(self) -> None: """Setup function for the JaxMLP class. From 51f7a073d66a49b9168d749db569f0f5241e5a06 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:43:26 -0400 Subject: [PATCH 15/41] Use TypeError instead of ValueError --- src/lanfactory/trainers/jax_mlp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 20c48ce..8cf81e6 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -50,7 +50,7 @@ def JaxMLPFactory( elif isinstance(network_config, dict): network_config_internal = network_config else: - raise ValueError( + raise TypeError( "network_config argument is not passed as either a dictionary or a string (path to a file)!" ) From 9ee816d7a1026aa553ff116154a2dfa318414a50 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:43:50 -0400 Subject: [PATCH 16/41] Simplify logic --- src/lanfactory/trainers/jax_mlp.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 8cf81e6..13dd126 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -120,20 +120,11 @@ def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray: for i, lyr in enumerate(self.layers): x = lyr(x) - if i != (len(self.layers) - 1): + if i != (len(self.layers) - 1) or self.activations[i] != "linear": x = self.activation_funs[i](x) - else: - if self.activations[i] == "linear": - pass - else: - x = self.activation_funs[i](x) - if (not self.train) and (self.train_output_type == "logprob"): - x = x # just for pedagogy - elif (not self.train) and (self.train_output_type == "logits"): + if (not self.train) and (self.train_output_type == "logits"): x = -jnp.log(1 + jnp.exp(-x)) - elif not self.train: # pragma: no cover - x = x # just for pedagogy return x From 273b6e9823f7da878e983e2c1606d022f80ee635 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:44:40 -0400 Subject: [PATCH 17/41] Simplify serialization --- src/lanfactory/trainers/jax_mlp.py | 31 ++++++++++++++---------------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 13dd126..2121b10 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -593,25 +593,22 @@ def train_and_evaluate( # Write to file train_state_path = f"{full_path}_train_state.jax" - file = open(train_state_path, "wb") - file.write(byte_output) - file.close() + Path(train_state_path).write_bytes(byte_output) print("Saving model parameters to: " + train_state_path) - config_path = f"{full_path}_train_config.pickle" - pickle.dump(self.train_config, open(config_path, "wb")) - print("Saving training config to: " + config_path) - - data_details_path = f"{full_path}_data_details.pickle" - pickle.dump( - { - "train_data_generator_config": self.train_dl.dataset.data_generator_config, - "train_data_file_ids": self.train_dl.dataset.file_ids, - "valid_data_generator_config": self.valid_dl.dataset.data_generator_config, - "valid_data_file_ids": self.valid_dl.dataset.file_ids, - }, - open(data_details_path, "wb"), - ) + config_path = Path(f"{full_path}_train_config.pickle") + config_path.write_bytes(pickle.dumps(self.train_config)) + print(f"Saving training config to: {config_path}") + + data_details_path = Path(f"{full_path}_data_details.pickle") + data_details = { + "train_data_generator_config": self.train_dl.dataset.data_generator_config, + "train_data_file_ids": self.train_dl.dataset.file_ids, + "valid_data_generator_config": self.valid_dl.dataset.data_generator_config, + "valid_data_file_ids": self.valid_dl.dataset.file_ids, + } + + data_details_path.write_bytes(pickle.dumps(data_details)) print("Saving training data details to: " + data_details_path) if self.mlflow_on: From ba656f0eae78c6ac68ea54fe4eb37249097d58c0 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:47:05 -0400 Subject: [PATCH 18/41] Remove unused vars --- src/lanfactory/trainers/jax_mlp.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 2121b10..3c16293 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -323,10 +323,10 @@ def loss_fn(params: dict) -> tuple[float, jnp.ndarray]: if train: grad_fn = jax.value_and_grad(loss_fn, has_aux=True) - (loss, pred), grads = grad_fn(state.params) + loss, _, grads = grad_fn(state.params) return grads, loss else: - loss, pred = loss_fn(state.params) + loss, _ = loss_fn(state.params) return loss return apply_model_core @@ -419,7 +419,7 @@ def run_epoch( # Run training for one epoch start_time = time() step = 0 - for X, y in tmp_dataloader: + for step, X, y in enumerate(tmp_dataloader): X_jax = jnp.array(X) y_jax = jnp.array(y) @@ -464,8 +464,6 @@ def run_epoch( else: pass - step += 1 - end_time = time() print( "Epoch " From 2b86808a38c99ea405cc13a6b041912ee9d4070c Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:47:16 -0400 Subject: [PATCH 19/41] Use TypeError instead of ValueError --- src/lanfactory/trainers/jax_mlp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 3c16293..313bf7b 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -208,7 +208,7 @@ def make_forward_partial( elif isinstance(state, dict): loaded_state = state else: - raise ValueError("state argument has to be a dictionary or a string!") + raise TypeError("state argument has to be a dictionary or a string!") # Make forward pass net_forward = partial(self.apply, loaded_state) From ef9ca8ea2d565b6c05da15e9f66c9d81b32ebce6 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:47:33 -0400 Subject: [PATCH 20/41] Improve exception --- src/lanfactory/trainers/jax_mlp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 313bf7b..bb44544 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -447,8 +447,8 @@ def run_epoch( if self.mlflow_on: try: mlflow.log_metric("loss", float(loss), step=int(state.step)) - except Exception: - pass + except Exception as e: + print(f"Failed to log metric to MLflow: {e}") elif verbose == 1: # pragma: no cover if (step % 1000) == 0: From 8ffb6eb837c1f7902b13992c0d2fdf25482211f3 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:47:44 -0400 Subject: [PATCH 21/41] Simplify logic --- src/lanfactory/trainers/jax_mlp.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index bb44544..219f3fa 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -450,19 +450,16 @@ def run_epoch( except Exception as e: print(f"Failed to log metric to MLflow: {e}") - elif verbose == 1: # pragma: no cover - if (step % 1000) == 0: - print( - train_str - + " - Step: " - + str(step) - + " of " - + str(cnt_max) - + " - Loss: " - + str(loss) - ) - else: - pass + elif verbose == 1 and (step % 1000) == 0: # pragma: no cover + print( + train_str + + " - Step: " + + str(step) + + " of " + + str(cnt_max) + + " - Loss: " + + str(loss) + ) end_time = time() print( From 2f43ab060b5fa3e630bb7b9b7fb2dbee3fd1980b Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:47:53 -0400 Subject: [PATCH 22/41] Use TypeError instead of ValueError --- src/lanfactory/trainers/jax_mlp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 219f3fa..d29c02d 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -531,7 +531,7 @@ def train_and_evaluate( # Initialize network if not isinstance(self.seed, int): - raise ValueError( + raise TypeError( "seed argument is not an integer, " + "please specify a valid seed to make this code reproducible!" ) From f1114a2c740f709c271ae0e0960f01629b8976a6 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:48:16 -0400 Subject: [PATCH 23/41] Add ruff lint exceptions --- pyproject.toml | 8 ++++++++ src/lanfactory/utils/util_funs.py | 0 2 files changed, 8 insertions(+) mode change 100755 => 100644 src/lanfactory/utils/util_funs.py diff --git a/pyproject.toml b/pyproject.toml index 7277bf6..045b4f3 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -120,6 +120,14 @@ line-length = 88 src = ["src/lanfactory", "tests"] exclude = ["notebooks/*", "docs/*", "docs/basic_tutorial/basic_tutorial.ipynb"] +[tool.ruff.lint] +ignore = [ + "B008", # Function calls in default arguments (Typer/FastAPI options) + "EXE001", # Shebang present but file is not executable + "EXE002", # Executable file fails to declare a shebang + "BLE001", # +] + # Coverage settings [tool.coverage.run] source = ["src"] diff --git a/src/lanfactory/utils/util_funs.py b/src/lanfactory/utils/util_funs.py old mode 100755 new mode 100644 From c1a1dca38b33969341db6ee96eeaec34901853da Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:57:24 -0400 Subject: [PATCH 24/41] Sort exported names --- src/lanfactory/trainers/__init__.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/lanfactory/trainers/__init__.py b/src/lanfactory/trainers/__init__.py index 6ba03f6..50627b4 100755 --- a/src/lanfactory/trainers/__init__.py +++ b/src/lanfactory/trainers/__init__.py @@ -15,18 +15,18 @@ __all__ = [ # Dataset and DataLoader helpers "DatasetTorch", - "make_dataloader", - "make_train_valid_dataloaders", - # Torch MLP - "TorchMLP", - "TorchMLPFactory", - "ModelTrainerTorchMLP", - "LoadTorchMLP", - "LoadTorchMLPInfer", + "JaxMLP", # Jax MLP "JaxMLPFactory", - "JaxMLP", + "LoadTorchMLP", + "LoadTorchMLPInfer", "ModelTrainerJaxMLP", + "ModelTrainerTorchMLP", + # Torch MLP + "TorchMLP", + "TorchMLPFactory", + "make_dataloader", + "make_train_valid_dataloaders", ] _DEPRECATED_ALIASES = { From 5bee526112f51dc4c9486b728c1432cbe902978b Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:57:39 -0400 Subject: [PATCH 25/41] Update torch train --- src/lanfactory/cli/torch_train.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/src/lanfactory/cli/torch_train.py b/src/lanfactory/cli/torch_train.py index 33eb128..c82bfb6 100644 --- a/src/lanfactory/cli/torch_train.py +++ b/src/lanfactory/cli/torch_train.py @@ -317,7 +317,7 @@ def main( # Mode 2: Validation - verify MLflow files exist in training_data_folder if mlflow_lineage_info and training_data_folder: expected_files = set(mlflow_lineage_info["all_files"]) - actual_files = set(f.name for f in valid_file_list) + actual_files = {f.name for f in valid_file_list} missing_files = expected_files - actual_files extra_files = actual_files - expected_files @@ -468,13 +468,8 @@ def main( ] ) - pickle.dump( - network_config, - open( - networks_path / file_name_suffix, - "wb", - ), - ) + file_path = networks_path / file_name_suffix + file_path.write_bytes(pickle.dumps(network_config)) # Load network net = lanfactory.trainers.TorchMLP( From 41e61c52ac3cc03c38ce2934d01f13681115eaf4 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:58:26 -0400 Subject: [PATCH 26/41] Make more readable --- src/lanfactory/hf/download.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/lanfactory/hf/download.py b/src/lanfactory/hf/download.py index 52bb17c..eb500de 100644 --- a/src/lanfactory/hf/download.py +++ b/src/lanfactory/hf/download.py @@ -123,10 +123,11 @@ def _download_model_hf( # pragma: no cover model_files = [f for f in all_files if f.startswith(path_prefix)] if not model_files: - raise FileNotFoundError( - f"No files found at {repo_id}/{path_prefix}. " - f"Available paths: {set(f.split('/')[0] for f in all_files if '/' in f)}" - ) + top_level_dirs = {f.split("/")[0] for f in all_files if "/" in f} + raise FileNotFoundError( + f"No files found at {repo_id}/{path_prefix}. " + f"Available paths: {top_level_dirs}" + ) if include_patterns: filtered_files = [] From e19327e528f7f3efd7ba59a72d498c312716b230 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:58:41 -0400 Subject: [PATCH 27/41] Update serialization --- src/lanfactory/cli/utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/lanfactory/cli/utils.py b/src/lanfactory/cli/utils.py index f3bad93..dc13f6b 100644 --- a/src/lanfactory/cli/utils.py +++ b/src/lanfactory/cli/utils.py @@ -89,9 +89,9 @@ def _make_train_network_configs( if save_name: save_folder = Path(save_folder) save_folder.mkdir(parents=True, exist_ok=True) # pragma: no cover - save_name = save_folder / save_name - pickle.dump(config_dict, open(save_name, "wb")) - print(f"Saved to: {save_name}") + save_path = save_folder / save_name + save_path.write_bytes(pickle.dumps(config_dict)) + print(f"Saved to: {save_path}") else: print("No save name provided, config not saved to file.") @@ -100,7 +100,7 @@ def _make_train_network_configs( def _get_train_network_config(yaml_config_path: str | Path | None = None, net_index=0): if yaml_config_path is not None: - basic_config = yaml.safe_load(open(yaml_config_path, "rb")) + basic_config = yaml.safe_load(Path(yaml_config_path).read_bytes()) network_type = basic_config["NETWORK_TYPE"] else: raise ValueError("No YAML config path provided") From 6c79703fc114f857e9c175cf6f4f043de2749b2e Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 14:59:29 -0400 Subject: [PATCH 28/41] Format file --- src/lanfactory/hf/download.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/lanfactory/hf/download.py b/src/lanfactory/hf/download.py index eb500de..ecef4e7 100644 --- a/src/lanfactory/hf/download.py +++ b/src/lanfactory/hf/download.py @@ -123,11 +123,11 @@ def _download_model_hf( # pragma: no cover model_files = [f for f in all_files if f.startswith(path_prefix)] if not model_files: - top_level_dirs = {f.split("/")[0] for f in all_files if "/" in f} - raise FileNotFoundError( - f"No files found at {repo_id}/{path_prefix}. " - f"Available paths: {top_level_dirs}" - ) + top_level_dirs = {f.split("/")[0] for f in all_files if "/" in f} + raise FileNotFoundError( + f"No files found at {repo_id}/{path_prefix}. " + f"Available paths: {top_level_dirs}" + ) if include_patterns: filtered_files = [] From 0cbcaa32f2a606ac5f34f426a3b5d15df566eaac Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 15:04:52 -0400 Subject: [PATCH 29/41] Fix unpacking --- src/lanfactory/trainers/jax_mlp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index d29c02d..f3e841f 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -419,7 +419,7 @@ def run_epoch( # Run training for one epoch start_time = time() step = 0 - for step, X, y in enumerate(tmp_dataloader): + for step, (X, y) in enumerate(tmp_dataloader): X_jax = jnp.array(X) y_jax = jnp.array(y) From 9f3170b00fcef26e49bb2c28a5193e5a1d7d8f63 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 15:48:25 -0400 Subject: [PATCH 30/41] Fix unpacking --- src/lanfactory/trainers/jax_mlp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index f3e841f..6c3e50f 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -323,7 +323,7 @@ def loss_fn(params: dict) -> tuple[float, jnp.ndarray]: if train: grad_fn = jax.value_and_grad(loss_fn, has_aux=True) - loss, _, grads = grad_fn(state.params) + (loss, _), grads = grad_fn(state.params) return grads, loss else: loss, _ = loss_fn(state.params) From 80ffaa55bb93aaeaa1d215e1c491748514f290b1 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 15:48:43 -0400 Subject: [PATCH 31/41] Fix logging message --- src/lanfactory/trainers/jax_mlp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 6c3e50f..bf508f5 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -604,7 +604,7 @@ def train_and_evaluate( } data_details_path.write_bytes(pickle.dumps(data_details)) - print("Saving training data details to: " + data_details_path) + print(f"Saving training data details to: {data_details_path}") if self.mlflow_on: try: From bf00bd466fdafed93339a2157266f5dd42f7d3a7 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 15:50:52 -0400 Subject: [PATCH 32/41] Deprecate mocking in test --- tests/test_cli_utils.py | 126 ++++++++++++++++++++-------------------- 1 file changed, 64 insertions(+), 62 deletions(-) diff --git a/tests/test_cli_utils.py b/tests/test_cli_utils.py index fe806ff..38644af 100644 --- a/tests/test_cli_utils.py +++ b/tests/test_cli_utils.py @@ -1,10 +1,10 @@ """Tests for CLI utilities.""" -from unittest.mock import patch, mock_open +import pickle import pytest - -from lanfactory.cli.utils import _make_train_network_configs, _get_train_network_config +import yaml +from lanfactory.cli.utils import _get_train_network_config, _make_train_network_configs def test_make_train_network_configs_with_dict_args(): @@ -44,25 +44,30 @@ def test_make_train_network_configs_without_save_name(): def test_make_train_network_configs_with_save_name(tmp_path): - """Test _make_train_network_configs with save_name (file written).""" + """Test _make_train_network_configs creates and writes the pickle file to disk.""" save_name = "test_config.pickle" + expected_file_path = tmp_path / save_name + + result = _make_train_network_configs( + training_data_folder="/fake/data", + train_val_split=0.9, + save_folder=str(tmp_path), + network_arg_dict=None, + train_arg_dict=None, + save_name=save_name, + ) + + assert result["config_file_name"] == expected_file_path.name - with patch("builtins.open", mock_open()), patch("pickle.dump") as mock_dump: - result = _make_train_network_configs( - training_data_folder="/fake/data", - train_val_split=0.9, - save_folder=str(tmp_path), - network_arg_dict=None, - train_arg_dict=None, - save_name=save_name, - ) + assert expected_file_path.is_file() - assert mock_dump.called - assert result["config_file_name"] == tmp_path / save_name + with open(expected_file_path, "rb") as f: + saved_data = pickle.load(f) + assert saved_data is not None -def test_get_train_network_config_lan(): - """Test _get_train_network_config with LAN network type.""" +def test_get_train_network_config_lan(tmp_path): + """Test _get_train_network_config with LAN network type using real file I/O.""" yaml_content = { "NETWORK_TYPE": "lan", "LAYER_SIZES": [[100, 100, 1]], @@ -83,20 +88,20 @@ def test_get_train_network_config_lan(): "MODEL": "ddm", } - with ( - patch("builtins.open", mock_open()), - patch("yaml.safe_load", return_value=yaml_content), - ): - result = _get_train_network_config(yaml_config_path="fake.yaml", net_index=0) + yaml_file = tmp_path / "test_config.yaml" + with open(yaml_file, "w") as f: + yaml.dump(yaml_content, f) - assert result["config_dict"]["network_config"]["train_output_type"] == "logprob" - assert result["config_dict"]["train_config"]["loss"] == "huber" - assert result["config_dict"]["train_config"]["features_key"] == "lan_data" - assert result["config_dict"]["train_config"]["label_key"] == "lan_labels" - assert result["extra_fields"]["model"] == "ddm" + result = _get_train_network_config(yaml_config_path=yaml_file, net_index=0) + assert result["config_dict"]["network_config"]["train_output_type"] == "logprob" + assert result["config_dict"]["train_config"]["loss"] == "huber" + assert result["config_dict"]["train_config"]["features_key"] == "lan_data" + assert result["config_dict"]["train_config"]["label_key"] == "lan_labels" + assert result["extra_fields"]["model"] == "ddm" -def test_get_train_network_config_cpn(): + +def test_get_train_network_config_cpn(tmp_path): """Test _get_train_network_config with CPN network type.""" yaml_content = { "NETWORK_TYPE": "cpn", @@ -118,19 +123,18 @@ def test_get_train_network_config_cpn(): "MODEL": "ddm", } - with ( - patch("builtins.open", mock_open()), - patch("yaml.safe_load", return_value=yaml_content), - ): - result = _get_train_network_config(yaml_config_path="fake.yaml", net_index=0) + config_file = tmp_path / "test_config.yaml" + config_file.write_text(yaml.dump(yaml_content)) + + result = _get_train_network_config(yaml_config_path=config_file, net_index=0) - assert result["config_dict"]["network_config"]["train_output_type"] == "logits" - assert result["config_dict"]["train_config"]["loss"] == "bcelogit" - assert result["config_dict"]["train_config"]["features_key"] == "cpn_data" - assert result["config_dict"]["train_config"]["label_key"] == "cpn_labels" + assert result["config_dict"]["network_config"]["train_output_type"] == "logits" + assert result["config_dict"]["train_config"]["loss"] == "bcelogit" + assert result["config_dict"]["train_config"]["features_key"] == "cpn_data" + assert result["config_dict"]["train_config"]["label_key"] == "cpn_labels" -def test_get_train_network_config_opn(): +def test_get_train_network_config_opn(tmp_path): """Test _get_train_network_config with OPN network type.""" yaml_content = { "NETWORK_TYPE": "opn", @@ -152,16 +156,15 @@ def test_get_train_network_config_opn(): "MODEL": "ddm", } - with ( - patch("builtins.open", mock_open()), - patch("yaml.safe_load", return_value=yaml_content), - ): - result = _get_train_network_config(yaml_config_path="fake.yaml", net_index=0) + config_file = tmp_path / "test_config_opn.yaml" + config_file.write_text(yaml.dump(yaml_content)) - assert result["config_dict"]["network_config"]["train_output_type"] == "logits" - assert result["config_dict"]["train_config"]["loss"] == "bcelogit" - assert result["config_dict"]["train_config"]["features_key"] == "opn_data" - assert result["config_dict"]["train_config"]["label_key"] == "opn_labels" + result = _get_train_network_config(yaml_config_path=config_file, net_index=0) + + assert result["config_dict"]["network_config"]["train_output_type"] == "logits" + assert result["config_dict"]["train_config"]["loss"] == "bcelogit" + assert result["config_dict"]["train_config"]["features_key"] == "opn_data" + assert result["config_dict"]["train_config"]["label_key"] == "opn_labels" def test_get_train_network_config_no_path(): @@ -170,7 +173,7 @@ def test_get_train_network_config_no_path(): _get_train_network_config(yaml_config_path=None) -def test_get_train_network_config_with_net_index(): +def test_get_train_network_config_with_net_index(tmp_path): """Test _get_train_network_config with different net_index.""" yaml_content = { "NETWORK_TYPE": "lan", @@ -192,17 +195,16 @@ def test_get_train_network_config_with_net_index(): "MODEL": "ddm", } - with ( - patch("builtins.open", mock_open()), - patch("yaml.safe_load", return_value=yaml_content), - ): - result = _get_train_network_config(yaml_config_path="fake.yaml", net_index=1) - - # layer_sizes comes directly from YAML (not modified) - assert result["config_dict"]["network_config"]["layer_sizes"] == [120, 120, 1] - # activations has output layer activation appended - assert result["config_dict"]["network_config"]["activations"] == [ - "relu", - "relu", - "linear", - ] + config_file = tmp_path / "test_config_net_index.yaml" + config_file.write_text(yaml.dump(yaml_content)) + + result = _get_train_network_config(yaml_config_path=config_file, net_index=1) + + # layer_sizes comes directly from YAML (not modified) + assert result["config_dict"]["network_config"]["layer_sizes"] == [120, 120, 1] + # activations has output layer activation appended + assert result["config_dict"]["network_config"]["activations"] == [ + "relu", + "relu", + "linear", + ] From 41598ef292764d0d2ef5ed2120dab0871ca77e7b Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 16:14:15 -0400 Subject: [PATCH 33/41] Update test jax mlp --- tests/test_jax_mlp.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/tests/test_jax_mlp.py b/tests/test_jax_mlp.py index 895a1e6..163c2bb 100644 --- a/tests/test_jax_mlp.py +++ b/tests/test_jax_mlp.py @@ -2,11 +2,10 @@ import pickle -import pytest import jax import jax.numpy as jnp - -from lanfactory.trainers.jax_mlp import JaxMLPFactory, JaxMLP +import pytest +from lanfactory.trainers.jax_mlp import JaxMLP, JaxMLPFactory def test_mlp_jax_factory_with_dict(): @@ -44,9 +43,9 @@ def test_mlp_jax_factory_with_string_path(tmp_path): assert model.layer_sizes == [100, 100, 1] -def test_mlp_jax_factory_raises_value_error(): +def test_mlp_jax_factory_raises_type_error(): """Test JaxMLPFactory raises ValueError for invalid network_config type.""" - with pytest.raises(ValueError, match="network_config argument is not passed"): + with pytest.raises(TypeError, match="network_config argument is not passed"): JaxMLPFactory(network_config=123, train=True) # Invalid type @@ -128,7 +127,7 @@ def test_mlp_jax_forward_with_non_linear_output_activation(): x = jax.random.normal(key, (5, 5)) # Initialize model state - key1, key2 = jax.random.split(key) + key1, _ = jax.random.split(key) state = model.init(key1, x) # Forward pass @@ -152,7 +151,7 @@ def test_mlp_jax_inference_mode_with_logits(): x = jax.random.normal(key, (5, 5)) # Initialize model state - key1, key2 = jax.random.split(key) + key1, _ = jax.random.split(key) state = model.init(key1, x) # Forward pass in inference mode @@ -163,8 +162,8 @@ def test_mlp_jax_inference_mode_with_logits(): def test_mlp_jax_load_state_from_file_error(): """Test JaxMLP load_state_from_file raises error when file_path is None.""" - from lanfactory.trainers.jax_mlp import JaxMLPFactory import pytest + from lanfactory.trainers.jax_mlp import JaxMLPFactory network_config = { "layer_sizes": [10, 10, 1], @@ -180,8 +179,8 @@ def test_mlp_jax_load_state_from_file_error(): def test_mlp_jax_load_state_from_file_without_input_dim(tmp_path): """Test JaxMLP load_state_from_file without providing input_dim.""" - from lanfactory.trainers.jax_mlp import JaxMLPFactory import flax.serialization + from lanfactory.trainers.jax_mlp import JaxMLPFactory network_config = { "layer_sizes": [10, 10, 1], @@ -242,8 +241,8 @@ def test_mlp_jax_make_forward_partial_with_dict_state(tmp_path): def test_mlp_jax_make_forward_partial_without_jit(tmp_path): """Test JaxMLP make_forward_partial without JIT compilation.""" - from lanfactory.trainers.jax_mlp import JaxMLPFactory import flax.serialization + from lanfactory.trainers.jax_mlp import JaxMLPFactory network_config = { "layer_sizes": [10, 10, 1], @@ -276,8 +275,8 @@ def test_mlp_jax_make_forward_partial_without_jit(tmp_path): def test_mlp_jax_make_forward_partial_invalid_state_type(): """Test JaxMLP make_forward_partial raises error with invalid state type.""" - from lanfactory.trainers.jax_mlp import JaxMLPFactory import pytest + from lanfactory.trainers.jax_mlp import JaxMLPFactory network_config = { "layer_sizes": [10, 10, 1], @@ -289,7 +288,7 @@ def test_mlp_jax_make_forward_partial_invalid_state_type(): # Test with invalid state type (list instead of dict or string) with pytest.raises( - ValueError, match="state argument has to be a dictionary or a string" + TypeError, match="state argument has to be a dictionary or a string" ): model.make_forward_partial( seed=42, input_dim=5, state=[1, 2, 3], add_jitted=True From f67a20d7ec150d0c48a9b0f3f0bd89077db15d11 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 16:34:26 -0400 Subject: [PATCH 34/41] Use tmp_path --- tests/conftest.py | 13 +++++-------- tests/constants.py | 1 - 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index ae0ae3e..fb3b2a5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -133,7 +133,7 @@ def _select_model( @pytest.fixture -def dummy_generator_config(model_selector): +def dummy_generator_config(model_selector, tmp_path): """Fixture providing a dummy model config for testing.""" def _dummy_generator_config(mode="random"): @@ -148,9 +148,7 @@ def _dummy_generator_config(mode="random"): generator_config["simulator"]["n_samples"] = ( TEST_GENERATOR_CONSTANTS.N_SAMPLES ) - generator_config["output"]["folder"] = os.path.join( - TEST_GENERATOR_CONSTANTS.OUT_FOLDER, str(uuid.uuid4()) - ) + generator_config["output"]["folder"] = str(tmp_path / str(uuid.uuid4())) generator_config["training"]["n_samples_per_param"] = ( TEST_GENERATOR_CONSTANTS.N_SAMPLES_BY_PARAMETER_SET ) @@ -169,7 +167,7 @@ def _dummy_generator_config(mode="random"): @pytest.fixture -def dummy_generator_config_simple_two_choices(model_selector): +def dummy_generator_config_simple_two_choices(model_selector, tmp_path): """Fixture providing a dummy model config for testing.""" def _dummy_generator_config_simple_two_choices(mode="random"): @@ -185,9 +183,8 @@ def _dummy_generator_config_simple_two_choices(mode="random"): generator_config["simulator"]["n_samples"] = ( TEST_GENERATOR_CONSTANTS.N_SAMPLES ) - generator_config["output"]["folder"] = os.path.join( - TEST_GENERATOR_CONSTANTS.OUT_FOLDER, str(uuid.uuid4()) - ) + generator_config["output"]["folder"] = str(tmp_path / str(uuid.uuid4())) + generator_config["training"]["n_samples_per_param"] = ( TEST_GENERATOR_CONSTANTS.N_SAMPLES_BY_PARAMETER_SET ) diff --git a/tests/constants.py b/tests/constants.py index dd65947..6404efa 100644 --- a/tests/constants.py +++ b/tests/constants.py @@ -31,7 +31,6 @@ class TestGeneratorConstants: N_TRAINING_SAMPLES: int = 2000 N_SAMPLES_BY_PARAMETER_SET: int = 2000 TEST_FOLDER: str = "tests/test_data" - OUT_FOLDER: str = "tests/test_data/lan_mlp/training_data" N_DATA_FILES: int = 2 DEVICE: str = "cpu" From c7ac744cefbc7d15f3951c9a4f27d84f891f8655 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 16:35:27 -0400 Subject: [PATCH 35/41] Organize imports --- tests/test_end_to_end_jax.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_end_to_end_jax.py b/tests/test_end_to_end_jax.py index 9a5a65d..dd84904 100644 --- a/tests/test_end_to_end_jax.py +++ b/tests/test_end_to_end_jax.py @@ -1,19 +1,19 @@ -import pytest -import ssms -import lanfactory +import logging import os -import numpy as np from copy import deepcopy + import jax.numpy as jnp +import lanfactory +import numpy as np +import pytest +import ssms import torch + from .constants import ( TEST_GENERATOR_CONSTANTS, TEST_MODEL_FOLDER_CONSTANTS_JAX, ) -# import logger -import logging - logger = logging.getLogger(__name__) LEN_FORWARD_PASS_DUMMY = 2000 From a0157e69643d818bd14884876b21c301267f4412 Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 16:35:39 -0400 Subject: [PATCH 36/41] Ignore .vscode dir --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index aa3bad6..df0292f 100755 --- a/.gitignore +++ b/.gitignore @@ -157,3 +157,4 @@ explorations/ __marimo__/ uv.lock +.vscode/ From 9301939e8fed95ee966bf41fe71bb68bd620d7fd Mon Sep 17 00:00:00 2001 From: cpaniaguam Date: Fri, 31 Jul 2026 17:10:34 -0400 Subject: [PATCH 37/41] Pass python-version to install uv action --- .github/workflows/run_tests.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index d97d32a..8240bdc 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -23,6 +23,7 @@ jobs: uses: astral-sh/setup-uv@v7 with: version: "0.12.0" + python-version: ${{ matrix.python-version }} # enable-cache: true # cache-dependency-glob: "pyproject.toml pdm.lock" From aff025176e538be8905c35a38d58db916352cb27 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Sat, 1 Aug 2026 21:56:17 +0000 Subject: [PATCH 38/41] chore(deps): bump actions/setup-python from 6 to 7 Bumps [actions/setup-python](https://github.com/actions/setup-python) from 6 to 7. - [Release notes](https://github.com/actions/setup-python/releases) - [Commits](https://github.com/actions/setup-python/compare/v6...v7) --- updated-dependencies: - dependency-name: actions/setup-python dependency-version: '7' dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] --- .github/workflows/run_tests.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index 23e1d2a..d482b6c 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -20,7 +20,7 @@ jobs: persist-credentials: false - name: Setup Python - uses: actions/setup-python@v6 + uses: actions/setup-python@v7 with: python-version: ${{ matrix.python-version }} @@ -68,7 +68,7 @@ jobs: persist-credentials: false - name: Setup Python - uses: actions/setup-python@v6 + uses: actions/setup-python@v7 with: python-version: "3.12" From 47af75d235fb5d1b856609de419213ead809b19a Mon Sep 17 00:00:00 2001 From: Carlos Paniagua Date: Mon, 3 Aug 2026 12:07:04 -0400 Subject: [PATCH 39/41] Fix linting workflow commands to use 'uv run' instead of 'uvx' --- .github/workflows/linting_formatting.yml | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/.github/workflows/linting_formatting.yml b/.github/workflows/linting_formatting.yml index bed620a..7235de2 100644 --- a/.github/workflows/linting_formatting.yml +++ b/.github/workflows/linting_formatting.yml @@ -17,8 +17,11 @@ jobs: with: version: "0.12.0" + - name: Install package + run: uv sync --group dev + - name: Check styling - run: uvx ruff format --check . + run: uv run ruff format --check . - name: Check linting - run: uvx ruff check src/lanfactory + run: uv run ruff check src/lanfactory From bbaa34d88ec1e33aa56d62257622977ff1ff5773 Mon Sep 17 00:00:00 2001 From: Carlos Paniagua Date: Tue, 4 Aug 2026 16:11:17 -0400 Subject: [PATCH 40/41] Fix linting --- pyproject.toml | 4 +-- tests/conftest.py | 1 - tests/hf/test_download.py | 1 - tests/hf/test_model_card.py | 1 - tests/hf/test_upload.py | 1 - tests/test_bayesflow_hssm_integration.py | 28 ++++++++-------- tests/test_bayesflow_nle_export.py | 39 +++++++++++++---------- tests/test_bayesflow_nre_export.py | 36 +++++++++++---------- tests/test_end_to_end_torch.py | 13 ++++---- tests/test_mlflow_integration.py | 18 +++++------ tests/test_network_inspectors_api.py | 1 - tests/test_network_inspectors_plotting.py | 1 - tests/test_sbi_embeddings.py | 27 ++++++++-------- tests/test_sbi_hssm_integration.py | 15 ++++----- tests/test_sbi_nle_export.py | 26 +++++++-------- tests/test_sbi_nre_export.py | 22 ++++++------- tests/test_sbi_spike_maf_roundtrip.py | 3 +- tests/test_sbi_spike_mlp_roundtrip.py | 3 +- tests/test_torch_mlp.py | 36 ++++++++++----------- tests/test_transform_onnx.py | 35 ++++++++++---------- tests/utils.py | 11 ++++--- 21 files changed, 162 insertions(+), 160 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 045b4f3..79f16d2 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,7 +82,7 @@ dev = [ "pytest-timer>=1.0.0", "pytest-xdist>=3.6.1", "pytest>=8.3.1", - "ruff>=0.14.4", + "ruff>=0.15.1", "types-PyYAML", "mlflow>=3.14.0", "jaxonnxruntime>=0.3", @@ -125,7 +125,7 @@ ignore = [ "B008", # Function calls in default arguments (Typer/FastAPI options) "EXE001", # Shebang present but file is not executable "EXE002", # Executable file fails to declare a shebang - "BLE001", # + "BLE001", # ] # Coverage settings diff --git a/tests/conftest.py b/tests/conftest.py index fb3b2a5..f99b1bd 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,5 @@ import logging import multiprocessing -import os import random import time import uuid diff --git a/tests/hf/test_download.py b/tests/hf/test_download.py index 686c185..47bee5b 100644 --- a/tests/hf/test_download.py +++ b/tests/hf/test_download.py @@ -3,7 +3,6 @@ from unittest.mock import patch import pytest - from lanfactory.hf.download import ( DEFAULT_REPO_ID, download_model, diff --git a/tests/hf/test_model_card.py b/tests/hf/test_model_card.py index 82e3525..c9a02a4 100644 --- a/tests/hf/test_model_card.py +++ b/tests/hf/test_model_card.py @@ -4,7 +4,6 @@ import pytest import yaml - from lanfactory.hf.model_card import ( ModelCardConfig, generate_readme, diff --git a/tests/hf/test_upload.py b/tests/hf/test_upload.py index 6eea1f6..fda6272 100644 --- a/tests/hf/test_upload.py +++ b/tests/hf/test_upload.py @@ -4,7 +4,6 @@ import pytest import yaml - from lanfactory.hf.upload import ( DEFAULT_INCLUDE_PATTERNS, DEFAULT_REPO_ID, diff --git a/tests/test_bayesflow_hssm_integration.py b/tests/test_bayesflow_hssm_integration.py index 066ba2c..4405882 100644 --- a/tests/test_bayesflow_hssm_integration.py +++ b/tests/test_bayesflow_hssm_integration.py @@ -24,23 +24,21 @@ os.environ["KERAS_BACKEND"] = "torch" os.environ.setdefault("KERAS_TORCH_DEVICE", "cpu") -from pathlib import Path # noqa: E402 - -import pytest # noqa: E402 +from pathlib import Path + +import bayesflow as bf +import keras +import numpy as np +import pandas as pd +import pytest +from bayesflow.datasets import OfflineDataset +from bayesflow.networks.inference.coupling.transforms import ( + AffineTransform, +) +from lanfactory.onnx import transform_bayesflow_to_onnx +from ssms.basic_simulators.simulator import simulator hssm = pytest.importorskip("hssm") - -import numpy as np # noqa: E402 -import pandas as pd # noqa: E402 - -import bayesflow as bf # noqa: E402 -import keras # noqa: E402 -from bayesflow.datasets import OfflineDataset # noqa: E402 -from bayesflow.networks.inference.coupling.transforms import AffineTransform # noqa: E402 -from ssms.basic_simulators.simulator import simulator # noqa: E402 - -from lanfactory.onnx import transform_bayesflow_to_onnx # noqa: E402 - _DDM_PARAM_NAMES = ["v", "a", "z", "t"] _DDM_PARAM_LOW = np.array([-2.0, 0.6, 0.3, 0.1], dtype=np.float32) _DDM_PARAM_HIGH = np.array([2.0, 1.8, 0.7, 0.5], dtype=np.float32) diff --git a/tests/test_bayesflow_nle_export.py b/tests/test_bayesflow_nle_export.py index ede588c..be92fb8 100644 --- a/tests/test_bayesflow_nle_export.py +++ b/tests/test_bayesflow_nle_export.py @@ -17,29 +17,32 @@ os.environ["KERAS_BACKEND"] = "torch" os.environ.setdefault("KERAS_TORCH_DEVICE", "cpu") -from pathlib import Path # noqa: E402 +from pathlib import Path +from typing import ClassVar + +import bayesflow as bf +import jax +import jax.numpy as jnp +import keras +import numpy as np +import onnxruntime as ort +import pytest +import torch +from bayesflow.datasets import OfflineDataset +from bayesflow.networks.inference.coupling.transforms import ( + AffineTransform, +) +from jaxonnxruntime import call_onnx, config +from lanfactory.onnx import transform_bayesflow_to_onnx -import jax # noqa: E402 +import onnx +from tests._onnx_utils import max_int64_abs # Same reason as the sbi test: ONNX shape/index tensors are int64; JAX's default # int32 silently truncates them inside jaxonnxruntime translation. jax.config.update("jax_enable_x64", True) -import jax.numpy as jnp # noqa: E402 -import numpy as np # noqa: E402 -import onnx # noqa: E402 -import onnxruntime as ort # noqa: E402 -import pytest # noqa: E402 -import torch # noqa: E402 -from jaxonnxruntime import call_onnx, config # noqa: E402 - -import bayesflow as bf # noqa: E402 -import keras # noqa: E402 -from bayesflow.datasets import OfflineDataset # noqa: E402 -from bayesflow.networks.inference.coupling.transforms import AffineTransform # noqa: E402 -from lanfactory.onnx import transform_bayesflow_to_onnx # noqa: E402 -from tests._onnx_utils import max_int64_abs # noqa: E402 # bayesflow under KERAS_BACKEND=torch globally disables autograd at import to # avoid excessive memory in long training loops. Restore the global default so @@ -400,7 +403,9 @@ def log_prob(self, samples, conditions): return (samples.sum() + conditions.sum()).reshape(1) class _Standardizer: - standardize_layers = {"inference_conditions": _FakeStandardizeLayer(_THETA_DIM)} + standardize_layers: ClassVar[dict[str, _FakeStandardizeLayer]] = { + "inference_conditions": _FakeStandardizeLayer(_THETA_DIM) + } class _Approx: adapter = None diff --git a/tests/test_bayesflow_nre_export.py b/tests/test_bayesflow_nre_export.py index b5258ab..ad014b1 100644 --- a/tests/test_bayesflow_nre_export.py +++ b/tests/test_bayesflow_nre_export.py @@ -14,25 +14,25 @@ os.environ["KERAS_BACKEND"] = "torch" os.environ.setdefault("KERAS_TORCH_DEVICE", "cpu") -from pathlib import Path # noqa: E402 - -import jax # noqa: E402 +from pathlib import Path +from typing import ClassVar + +import bayesflow as bf +import jax +import jax.numpy as jnp +import keras +import numpy as np +import onnxruntime as ort +import pytest +import torch +from bayesflow.datasets import OfflineDataset +from jaxonnxruntime import call_onnx, config +from lanfactory.onnx import transform_bayesflow_to_onnx + +import onnx jax.config.update("jax_enable_x64", True) -import jax.numpy as jnp # noqa: E402 -import numpy as np # noqa: E402 -import onnx # noqa: E402 -import onnxruntime as ort # noqa: E402 -import pytest # noqa: E402 -import torch # noqa: E402 -from jaxonnxruntime import call_onnx, config # noqa: E402 - -import bayesflow as bf # noqa: E402 -import keras # noqa: E402 -from bayesflow.datasets import OfflineDataset # noqa: E402 - -from lanfactory.onnx import transform_bayesflow_to_onnx # noqa: E402 # bayesflow under KERAS_BACKEND=torch globally disables autograd at import to # avoid excessive memory in long training loops. Restore the global default so @@ -261,7 +261,9 @@ def __call__(self, hidden): return hidden.sum().reshape(1, 1) class _Standardizer: - standardize_layers = {"inference_conditions": _FakeStandardizeLayer(_X_DIM)} + standardize_layers: ClassVar[dict[str, _FakeStandardizeLayer]] = { + "inference_conditions": _FakeStandardizeLayer(_X_DIM) + } class _Approx: def __init__(self): diff --git a/tests/test_end_to_end_torch.py b/tests/test_end_to_end_torch.py index 22a82ae..8a5c5c9 100644 --- a/tests/test_end_to_end_torch.py +++ b/tests/test_end_to_end_torch.py @@ -1,17 +1,18 @@ -import pytest -import ssms -import lanfactory +import logging import os -import numpy as np from copy import deepcopy + +import lanfactory +import numpy as np +import pytest +import ssms import torch + from .constants import ( TEST_GENERATOR_CONSTANTS, TEST_MODEL_FOLDER_CONSTANTS_TORCH, ) -import logging - logger = logging.getLogger(__name__) LEN_FORWARD_PASS_DUMMY = 2000 diff --git a/tests/test_mlflow_integration.py b/tests/test_mlflow_integration.py index c2f348a..7d3ed07 100644 --- a/tests/test_mlflow_integration.py +++ b/tests/test_mlflow_integration.py @@ -4,8 +4,8 @@ import shutil from pathlib import Path -import pytest import numpy as np +import pytest try: import mlflow @@ -51,7 +51,7 @@ def cleanup_mlflow(): try: mlflow.set_tracking_uri(original_uri) except Exception: - pass + mlflow.set_tracking_uri("file:./mlruns") @pytest.fixture @@ -306,9 +306,9 @@ def test_jax_trainer_mlflow_logging( tmp_path, ): """Test that JAX trainer logs to MLflow correctly.""" - from torch.utils.data import DataLoader from lanfactory.trainers.jax_mlp import JaxMLPFactory, ModelTrainerJaxMLP from lanfactory.trainers.torch_mlp import DatasetTorch + from torch.utils.data import DataLoader tracking_uri = test_mlflow_dir["tracking_uri"] artifact_location = test_mlflow_dir["artifact_location"] @@ -400,12 +400,12 @@ def test_pytorch_trainer_mlflow_logging( tmp_path, ): """Test that PyTorch trainer logs to MLflow correctly.""" - from torch.utils.data import DataLoader from lanfactory.trainers.torch_mlp import ( - TorchMLP, - ModelTrainerTorchMLP, DatasetTorch, + ModelTrainerTorchMLP, + TorchMLP, ) + from torch.utils.data import DataLoader tracking_uri = test_mlflow_dir["tracking_uri"] artifact_location = test_mlflow_dir["artifact_location"] @@ -639,12 +639,12 @@ def test_trainer_without_mlflow( tmp_path, ): """Test that trainers work correctly when MLflow is disabled.""" - from torch.utils.data import DataLoader from lanfactory.trainers.torch_mlp import ( - TorchMLP, - ModelTrainerTorchMLP, DatasetTorch, + ModelTrainerTorchMLP, + TorchMLP, ) + from torch.utils.data import DataLoader # Generate minimal training data gen_configs = dummy_generator_config_simple_two_choices() diff --git a/tests/test_network_inspectors_api.py b/tests/test_network_inspectors_api.py index 606eb95..fcd8c88 100644 --- a/tests/test_network_inspectors_api.py +++ b/tests/test_network_inspectors_api.py @@ -5,7 +5,6 @@ import numpy as np import pandas as pd import pytest - from lanfactory.network_inspectors import api diff --git a/tests/test_network_inspectors_plotting.py b/tests/test_network_inspectors_plotting.py index 4498ad0..5d35f98 100644 --- a/tests/test_network_inspectors_plotting.py +++ b/tests/test_network_inspectors_plotting.py @@ -5,7 +5,6 @@ import numpy as np import pandas as pd import plotly.graph_objects as go - from lanfactory.network_inspectors.config import ModelSpec, PlotConfig from lanfactory.network_inspectors.plotting import plot_manifold diff --git a/tests/test_sbi_embeddings.py b/tests/test_sbi_embeddings.py index b39a961..2c390fd 100644 --- a/tests/test_sbi_embeddings.py +++ b/tests/test_sbi_embeddings.py @@ -13,23 +13,22 @@ from pathlib import Path import jax +import numpy as np +import onnxruntime as ort +import pytest +import torch +from jaxonnxruntime import call_onnx, config +from lanfactory.onnx import transform_sbi_to_onnx +from sbi.inference import NRE_A +from sbi.neural_nets import classifier_nn +from sbi.neural_nets.embedding_nets import CNNEmbedding, FCEmbedding +from sbi.utils import BoxUniform +from torch import nn + +import onnx jax.config.update("jax_enable_x64", True) -import numpy as np # noqa: E402 -import onnx # noqa: E402 -import onnxruntime as ort # noqa: E402 -import pytest # noqa: E402 -import torch # noqa: E402 -from jaxonnxruntime import call_onnx, config # noqa: E402 -from sbi.inference import NRE_A # noqa: E402 -from sbi.neural_nets import classifier_nn # noqa: E402 -from sbi.neural_nets.embedding_nets import CNNEmbedding, FCEmbedding # noqa: E402 -from sbi.utils import BoxUniform # noqa: E402 -from torch import nn # noqa: E402 - -from lanfactory.onnx import transform_sbi_to_onnx # noqa: E402 - config.update("jaxort_only_allow_initializers_as_static_args", False) # sbi's build_mlp_classifier defaults to nn.LayerNorm between hidden layers, but diff --git a/tests/test_sbi_hssm_integration.py b/tests/test_sbi_hssm_integration.py index fcf34ba..97b41ab 100644 --- a/tests/test_sbi_hssm_integration.py +++ b/tests/test_sbi_hssm_integration.py @@ -26,14 +26,13 @@ # Skip cleanly when HSSM is not in the environment. hssm = pytest.importorskip("hssm") -import numpy as np # noqa: E402 -import pandas as pd # noqa: E402 -import torch # noqa: E402 -from sbi.inference import NLE_A # noqa: E402 -from sbi.utils import BoxUniform # noqa: E402 -from ssms.basic_simulators.simulator import simulator # noqa: E402 - -from lanfactory.onnx import transform_sbi_to_onnx # noqa: E402 +import numpy as np +import pandas as pd +import torch +from lanfactory.onnx import transform_sbi_to_onnx +from sbi.inference import NLE_A +from sbi.utils import BoxUniform +from ssms.basic_simulators.simulator import simulator # DDM parameter order matches sbi simulator inputs and HSSM defaults. _DDM_PARAM_NAMES = ["v", "a", "z", "t"] diff --git a/tests/test_sbi_nle_export.py b/tests/test_sbi_nle_export.py index d2a41fc..52bb772 100644 --- a/tests/test_sbi_nle_export.py +++ b/tests/test_sbi_nle_export.py @@ -17,18 +17,18 @@ # wrong numerical values (~0.5 drift from the torch reference on MAF log_prob). jax.config.update("jax_enable_x64", True) -import jax.numpy as jnp # noqa: E402 -import numpy as np # noqa: E402 -import onnx # noqa: E402 -import onnxruntime as ort # noqa: E402 -import pytest # noqa: E402 -import torch # noqa: E402 -from jaxonnxruntime import call_onnx, config # noqa: E402 -from sbi.inference import NLE_A # noqa: E402 -from sbi.utils import BoxUniform # noqa: E402 - -from lanfactory.onnx import transform_sbi_to_onnx # noqa: E402 -from tests._onnx_utils import max_int64_abs # noqa: E402 +import jax.numpy as jnp +import numpy as np +import onnxruntime as ort +import pytest +import torch +from jaxonnxruntime import call_onnx, config +from lanfactory.onnx import transform_sbi_to_onnx +from sbi.inference import NLE_A +from sbi.utils import BoxUniform + +import onnx +from tests._onnx_utils import max_int64_abs # Same friction as C2's MAF spike — torch.onnx.export emits Reshape shapes as # Constant nodes. HSSM's onnx2jax patch (commit 2e76516) sets this globally for @@ -191,7 +191,7 @@ def test_nle_log_prob_ordering_matches_analytical_gaussian( def test_transform_rejects_unsupported_score_estimator(tmp_path: Path) -> None: """Estimators in the unsupported set should fail loudly.""" - class ScoreEstimator(torch.nn.Module): # noqa: D401 - name is the signal + class ScoreEstimator(torch.nn.Module): pass with pytest.raises(ValueError, match="does not support"): diff --git a/tests/test_sbi_nre_export.py b/tests/test_sbi_nre_export.py index 20a7433..7ced473 100644 --- a/tests/test_sbi_nre_export.py +++ b/tests/test_sbi_nre_export.py @@ -14,17 +14,17 @@ # x64 required before any JAX import — see test_sbi_nle_export.py for details. jax.config.update("jax_enable_x64", True) -import jax.numpy as jnp # noqa: E402 -import numpy as np # noqa: E402 -import onnx # noqa: E402 -import onnxruntime as ort # noqa: E402 -import pytest # noqa: E402 -import torch # noqa: E402 -from jaxonnxruntime import call_onnx, config # noqa: E402 -from sbi.inference import NRE_A # noqa: E402 -from sbi.utils import BoxUniform # noqa: E402 - -from lanfactory.onnx import transform_sbi_to_onnx # noqa: E402 +import jax.numpy as jnp +import numpy as np +import onnxruntime as ort +import pytest +import torch +from jaxonnxruntime import call_onnx, config +from lanfactory.onnx import transform_sbi_to_onnx +from sbi.inference import NRE_A +from sbi.utils import BoxUniform + +import onnx config.update("jaxort_only_allow_initializers_as_static_args", False) diff --git a/tests/test_sbi_spike_maf_roundtrip.py b/tests/test_sbi_spike_maf_roundtrip.py index f0eb174..dec708e 100644 --- a/tests/test_sbi_spike_maf_roundtrip.py +++ b/tests/test_sbi_spike_maf_roundtrip.py @@ -11,7 +11,6 @@ import jax import numpy as np -import onnx import onnxruntime as ort import pytest import torch @@ -25,6 +24,8 @@ ) from torch import nn +import onnx + # Friction discovered in C2: nflows' MAF exports a Reshape whose shape argument # is a Constant node (not a model initializer). jaxonnxruntime's default strict # mode rejects this. The flag below tells jaxonnxruntime to treat Constant nodes diff --git a/tests/test_sbi_spike_mlp_roundtrip.py b/tests/test_sbi_spike_mlp_roundtrip.py index 100beb0..652de04 100644 --- a/tests/test_sbi_spike_mlp_roundtrip.py +++ b/tests/test_sbi_spike_mlp_roundtrip.py @@ -10,13 +10,14 @@ import jax import numpy as np -import onnx import onnxruntime as ort import pytest import torch from jaxonnxruntime import call_onnx from torch import nn +import onnx + @pytest.mark.flaky(reruns=2) def test_mlp_three_way_agreement(tmp_path: Path) -> None: diff --git a/tests/test_torch_mlp.py b/tests/test_torch_mlp.py index da38fe9..8076907 100644 --- a/tests/test_torch_mlp.py +++ b/tests/test_torch_mlp.py @@ -1,15 +1,15 @@ """Tests for the DatasetTorch class and related components.""" import pickle +from unittest.mock import MagicMock + import numpy as np import pytest import torch -from unittest.mock import MagicMock - from lanfactory.trainers.torch_mlp import ( DatasetTorch, - ModelTrainerTorchMLP, LoadTorchMLPInfer, + ModelTrainerTorchMLP, TorchMLP, ) @@ -110,7 +110,7 @@ def test_dataset_torch_getitem_loads_file_on_first_access( # First access should load file and populate tmp_data assert dataset.tmp_data == {} # Initially empty - X, y = dataset[0] + _, _ = dataset[0] # Now tmp_data should be populated assert dataset.tmp_data != {} @@ -132,9 +132,9 @@ def test_dataset_torch_getitem_multiple_batches_same_file( ) # Access multiple batches from the same file - X0, y0 = dataset[0] - X1, y1 = dataset[1] - X2, y2 = dataset[2] + X0, _y0 = dataset[0] + X1, _y1 = dataset[1] + X2, _y2 = dataset[2] # All should have correct shape assert X0.shape == (200, 6) @@ -163,22 +163,22 @@ def test_dataset_torch_getitem_crosses_file_boundary( # So indices 0-3 are from file 0, 4-7 from file 1, 8-11 from file 2 # Access batch from first file - X0, y0 = dataset[0] + _X0, _y0 = dataset[0] first_file_data = dataset.tmp_data["lan_data"].copy() # Access batch from same file - X3, y3 = dataset[3] + _X3, _y3 = dataset[3] assert np.array_equal(dataset.tmp_data["lan_data"], first_file_data) # Access first batch from second file - should trigger file load - X4, y4 = dataset[4] + _X4, _y4 = dataset[4] second_file_data = dataset.tmp_data["lan_data"] # Data should be different (new file loaded) assert not np.array_equal(second_file_data, first_file_data) # Access first batch from third file - X8, y8 = dataset[8] + _X8, _y8 = dataset[8] third_file_data = dataset.tmp_data["lan_data"] # Should be different from second file @@ -200,7 +200,7 @@ def test_dataset_torch_getitem_with_label_bounds( label_key="lan_labels", ) - X, y = dataset[0] + _, y = dataset[0] # Labels should be clipped to bounds assert np.all(y >= -10.0) @@ -271,7 +271,7 @@ def test_dataset_torch_empty_tmp_data_triggers_load( ) # Access batch 2 (not 0), but tmp_data is empty - X, y = dataset[2] + X, _y = dataset[2] # Should work - file was loaded due to empty tmp_data assert X.shape == (200, 6) @@ -342,7 +342,7 @@ def test_dataset_torch_label_bounds(tmp_path): label_key="lan_labels", ) - X, y = dataset[0] + _, y = dataset[0] # Labels should be clipped to bounds assert np.all(y >= -10.0) @@ -367,7 +367,7 @@ def test_dataset_torch_3d_labels_raises_error(tmp_path): ) with pytest.raises(ValueError, match="Label data has unexpected shape"): - X, y = dataset[0] + _X, _y = dataset[0] def test_dataset_torch_batch_size_not_divisible_raises_error(tmp_path): @@ -987,12 +987,12 @@ def test_torch_mlp_with_non_linear_output_activation(): def test_model_trainer_torch_mlp_with_none_train_config(create_mock_data_files): """Test ModelTrainerTorchMLP raises error when train_config is None.""" + import pytest from lanfactory.trainers.torch_mlp import ( DatasetTorch, ModelTrainerTorchMLP, TorchMLP, ) - import pytest file_list = create_mock_data_files(n_files=1) @@ -1113,7 +1113,7 @@ def test_make_train_valid_dataloaders_raises_no_valid_files(create_mock_data_fil def test_torch_mlp_factory_with_dict(): """Test TorchMLPFactory with dict config.""" - from lanfactory.trainers.torch_mlp import TorchMLPFactory, TorchMLP + from lanfactory.trainers.torch_mlp import TorchMLP, TorchMLPFactory network_config = { "layer_sizes": [10, 10, 1], @@ -1129,7 +1129,7 @@ def test_torch_mlp_factory_with_dict(): def test_torch_mlp_factory_with_pickle_path(tmp_path): """Test TorchMLPFactory with path to pickled config.""" - from lanfactory.trainers.torch_mlp import TorchMLPFactory, TorchMLP + from lanfactory.trainers.torch_mlp import TorchMLP, TorchMLPFactory network_config = { "layer_sizes": [10, 10, 1], diff --git a/tests/test_transform_onnx.py b/tests/test_transform_onnx.py index ec663e2..451c9fd 100644 --- a/tests/test_transform_onnx.py +++ b/tests/test_transform_onnx.py @@ -5,8 +5,7 @@ import pytest import torch - -from lanfactory.onnx.transform_onnx import transform_to_onnx, main +from lanfactory.onnx.transform_onnx import main, transform_to_onnx @pytest.fixture @@ -165,14 +164,16 @@ def test_transform_to_onnx_creates_correct_input_tensor(mock_network_config): def test_transform_to_onnx_missing_config_file(): """Test that transform_to_onnx raises error for missing config file.""" # Mock open to raise FileNotFoundError - with patch("builtins.open", side_effect=FileNotFoundError("File not found")): - with pytest.raises(FileNotFoundError): - transform_to_onnx( - network_config_file="/nonexistent/config.pickle", - state_dict_file="/fake/state.pt", - input_shape=6, - output_onnx_file="/fake/output.onnx", - ) + with ( + patch("builtins.open", side_effect=FileNotFoundError("File not found")), + pytest.raises(FileNotFoundError), + ): + transform_to_onnx( + network_config_file="/nonexistent/config.pickle", + state_dict_file="/fake/state.pt", + input_shape=6, + output_onnx_file="/fake/output.onnx", + ) def test_transform_to_onnx_missing_state_dict_file(mock_network_config): @@ -210,14 +211,14 @@ def test_transform_to_onnx_invalid_pickle_file(): "lanfactory.onnx.transform_onnx.pickle.load", side_effect=pickle.UnpicklingError("Invalid pickle"), ), + pytest.raises(pickle.UnpicklingError), ): - with pytest.raises(pickle.UnpicklingError): - transform_to_onnx( - network_config_file=config_file, - state_dict_file="/fake/state.pt", - input_shape=6, - output_onnx_file="/fake/output.onnx", - ) + transform_to_onnx( + network_config_file=config_file, + state_dict_file="/fake/state.pt", + input_shape=6, + output_onnx_file="/fake/output.onnx", + ) def test_main_calls_transform_to_onnx(): diff --git a/tests/utils.py b/tests/utils.py index 678a944..d177b0d 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,7 +1,6 @@ -import shutil -import pathlib - import logging +import pathlib +import shutil logger = logging.getLogger(__name__) @@ -52,7 +51,7 @@ def clean_out_folder(folder: str | pathlib.Path | None = None, dry_run=True) -> except PermissionError: logger.error(f"Permission denied when trying to remove folder '{folder}'.") except Exception as e: - logger.error(f"Error removing folder '{folder}': {str(e)}") + logger.error(f"Error removing folder '{folder}': {e!s}") else: logger.error(f"Folder '{folder}' does not exist.") @@ -61,7 +60,7 @@ def print_tree( path: pathlib.Path | str, prefix: str = "", logger: logging.Logger | None = logger, - out_str_list: list[str] = [], + out_str_list: list[str] | None = None, ) -> list[str]: """Print a directory tree structure starting from the given path. @@ -86,6 +85,8 @@ def print_tree( print_tree(Path("./my_directory")) """ path = pathlib.Path(path) + if out_str_list is None: + out_str_list = [] contents = sorted(path.iterdir(), key=lambda p: (not p.is_dir(), p.name.lower())) for index, item in enumerate(contents): connector = "└── " if index == len(contents) - 1 else "├── " From b7d0f7d69099c9d329bc1f2d161b5eedd4fcca13 Mon Sep 17 00:00:00 2001 From: Carlos Paniagua Date: Tue, 4 Aug 2026 16:13:36 -0400 Subject: [PATCH 41/41] Remove unnecessary blank line in test_bayesflow_nle_export.py --- tests/test_bayesflow_nle_export.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_bayesflow_nle_export.py b/tests/test_bayesflow_nle_export.py index be92fb8..2f27a76 100644 --- a/tests/test_bayesflow_nle_export.py +++ b/tests/test_bayesflow_nle_export.py @@ -43,7 +43,6 @@ jax.config.update("jax_enable_x64", True) - # bayesflow under KERAS_BACKEND=torch globally disables autograd at import to # avoid excessive memory in long training loops. Restore the global default so # that subsequent tests (e.g. test_sbi_*) that rely on autograd-by-default work.