l2co-tasks 0.1.0__tar.gz

l2co_tasks-0.1.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, Martin van der Schelling
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

l2co_tasks-0.1.0/MANIFEST.in

@@ -0,0 +1,3 @@
+include LICENSE
+include README.md
+recursive-include src/l2co_tasks/conf *.yaml

l2co_tasks-0.1.0/PKG-INFO

@@ -0,0 +1,138 @@
+Metadata-Version: 2.4
+Name: l2co-tasks
+Version: 0.1.0
+Summary: Optimization tasks compatible with the L2CO library
+Author-email: Martin van der Schelling <m.p.vanderschelling@tudelft.nl>
+Maintainer-email: Martin van der Schelling <m.p.vanderschelling@tudelft.nl>
+License-Expression: BSD-3-Clause
+Project-URL: bugs, https://github.com/bessagroup/l2co-tasks/issues
+Project-URL: homepage, https://github.com/bessagroup/l2co-tasks
+Keywords: optimization,tasks,benchmarking
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: bbob-jax>=1.6.4
+Requires-Dist: cloudpickle>=3.1.1
+Requires-Dist: equinox>=0.13.0
+Requires-Dist: f3dasm
+Requires-Dist: hydra-core>=1.3.2
+Requires-Dist: jax
+Requires-Dist: jaxtyping>=0.3.2
+Requires-Dist: mnist1d
+Requires-Dist: numpy
+Requires-Dist: omegaconf
+Requires-Dist: optax>=0.2.5
+Requires-Dist: pandas
+Provides-Extra: dev
+Requires-Dist: setuptools>=77.0.0; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: wheel; extra == "dev"
+Requires-Dist: pre-commit>=4.4.0; extra == "dev"
+Requires-Dist: ruff>=0.14.10; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: mkdocs>=1.6.0; extra == "docs"
+Requires-Dist: mkdocs-gen-files>=0.5.0; extra == "docs"
+Requires-Dist: mkdocs-include-exclude-files>=0.1.0; extra == "docs"
+Requires-Dist: mkdocs-ipynb>=0.1.1; extra == "docs"
+Requires-Dist: mkdocs-material>=9.6.20; extra == "docs"
+Requires-Dist: mkdocstrings>=0.30.1; extra == "docs"
+Requires-Dist: mkdocstrings-python>=1.18.2; extra == "docs"
+Requires-Dist: pymdown-extensions>=10.16.1; extra == "docs"
+Provides-Extra: tests
+Requires-Dist: pytest>=8.4.1; extra == "tests"
+Requires-Dist: pytest-cov>=6.2.1; extra == "tests"
+Dynamic: license-file
+
+![Bessa Research Group](img/bessa_group_logo.png)
+
+# L2CO Tasks
+
+| [**GitHub**](https://github.com/bessagroup/l2co-tasks)
+| [**Documentation**](https://l2co-tasks.readthedocs.io/en/latest/)
+
+Optimization tasks compatible with the L2CO library
+
+**First publication:** April 13, 2026
+
+***
+
+## Summary
+
+`l2co-tasks` provides a unified collection of optimization tasks for the [L2CO](https://github.com/bessagroup/l2co) ecosystem. Each `Task` bundles a JAX/Equinox model, a loss function, an optional dataset, and metadata tags into a single, serializable object that any L2CO optimizer or rollout can consume. The package ships ready-made task families — analytic black-box benchmarks ([BBOB](https://github.com/bessagroup/bbob-jax) and CEC 2005), quadratic problems, supervised-learning tasks (spiral, MNIST-1D, Gaussian classification), and PINN-style PDE problems (convection, reaction, wave) — together with `create_*_task` factories and `f3dasm` samplers for building experiment datasets.
+
+## Statement of need
+
+Research on learning to optimize and optimizer selection requires evaluating many optimizers across a diverse, well-characterized set of problems — but these problems usually come from incompatible sources with different input ranges, calling conventions, and metadata. `l2co-tasks` standardizes them behind a single `Task` abstraction: inputs are normalized to `[0, 1]^d` and scaled inside the loss, stochasticity and known global minima are tracked explicitly, and every task carries a hashable tag and serializes to a single `.eqx` file. This makes tasks reproducible, portable, and directly pluggable into the [`l2co`](https://github.com/bessagroup/l2co) / [`rl2co`](https://github.com/bessagroup/rl2co) rollout machinery and [`f3dasm`](https://github.com/bessagroup/f3dasm) experiment pipelines.
+
+## Authorship
+
+**Authors**:
+- Martin van der Schelling ([m.p.vanderschelling@tudelft.nl](mailto:m.p.vanderschelling@tudelft.nl))
+
+**Authors affiliation:**
+- Delft University of Technology (Bessa Research Group)
+
+**Maintainer:**
+- Martin van der Schelling ([m.p.vanderschelling@tudelft.nl](mailto:m.p.vanderschelling@tudelft.nl))
+
+**Maintainer affiliation:**
+- Delft University of Technology (Bessa Research Group)
+
+
+## Getting started
+
+`l2co-tasks` is `uv`-managed and depends on an editable install of a sibling [`f3dasm`](https://github.com/bessagroup/f3dasm) checkout, so lay the repositories out side-by-side before syncing:
+
+```bash
+git clone https://github.com/bessagroup/f3dasm.git
+git clone https://github.com/bessagroup/l2co-tasks.git
+cd l2co-tasks
+uv sync
+```
+
+Create a task from one of the factories and evaluate its loss:
+
+```python
+from l2co_tasks import create_bbob_task
+
+task = create_bbob_task(fn_name="sphere", seed=0, dimensionality=2)
+loss = task.loss_fn(task.model)   # model is the [0, 1]^d input vector
+```
+
+Tasks serialize to a single-file `.eqx` format via `Task.save` / `Task.load`. See the [documentation](https://l2co-tasks.readthedocs.io/en/latest/) for the full list of task families and `create_*_task` factories. To build your own task from scratch, see the [Create your own task](./docs/create_task.ipynb) guide.
+
+## Examples and benchmarks
+
+Two runnable notebooks demonstrate and benchmark the core functionality end-to-end. Both run against the installed package after `uv sync` and are rendered in the [documentation](https://l2co-tasks.readthedocs.io/en/latest/):
+
+- [Create your own task](./docs/create_task.ipynb) — builds a `Task` from scratch and exercises the `Task` API (loss evaluation, serialization, metadata).
+- [PINN benchmark tasks](./docs/pinn_tasks.ipynb) — instantiates and evaluates the physics-informed PDE benchmark suite (Helmholtz, Burgers, Euler, Stokes, PK–PD).
+
+## Community Support
+
+If you find any **issues, bugs or problems** with this package, please use the [GitHub issue tracker](https://github.com/bessagroup/l2co-tasks/issues) to report them.
+
+## License
+
+Copyright (c) 2026, Martin van der Schelling
+
+All rights reserved.
+
+This project is licensed under the BSD 3-Clause License. See [LICENSE](https://github.com/bessagroup/l2co-tasks/blob/main/LICENSE) for the full license text.
+
+## Related repositories
+
+This package is part of the L2CO ecosystem developed in the [Bessa Research Group](https://github.com/bessagroup). The repositories below work together:
+
+- [l2co](https://github.com/bessagroup/L2CO) — Learning to Choose Optimizers: a meta-learner that selects an optimizer from problem features before any evaluations, then reassesses that choice from the observed optimization trajectory.
+- [rl2co](https://github.com/bessagroup/rl2co) — Reinforcement Learning to Choose Optimizers: a JAX-based RL agent that dynamically switches between optimizers during a run.
+- [l2co-tasks](https://github.com/bessagroup/l2co-tasks) — Optimization task definitions (BBOB, CEC 2005, PDE, spiral, …) compatible with the L2CO library.
+- [l2co_experiments](https://github.com/bessagroup/l2co_experiments) — Hydra + f3dasm experiment pipelines (dataset creation, training, rollouts, figures) for the L2CO studies.
+- [agentic-l2co](https://github.com/bessagroup/agentic-l2co) — An LLM-agent drop-in replacement for `l2co.L2COModel`, driving two-stage optimizer selection with an Ollama-hosted LLM.
+- [bbob-jax](https://github.com/bessagroup/bbob-jax) — JAX implementations of the BBOB and CEC 2005 black-box optimization benchmark functions.
+- [f3dasm](https://github.com/bessagroup/f3dasm) — Framework for Data-Driven Design and Analysis of Structures and Materials; provides `ExperimentData`, pipelines, and SLURM orchestration.
+

l2co_tasks-0.1.0/README.md

@@ -0,0 +1,89 @@
+![Bessa Research Group](img/bessa_group_logo.png)
+
+# L2CO Tasks
+
+| [**GitHub**](https://github.com/bessagroup/l2co-tasks)
+| [**Documentation**](https://l2co-tasks.readthedocs.io/en/latest/)
+
+Optimization tasks compatible with the L2CO library
+
+**First publication:** April 13, 2026
+
+***
+
+## Summary
+
+`l2co-tasks` provides a unified collection of optimization tasks for the [L2CO](https://github.com/bessagroup/l2co) ecosystem. Each `Task` bundles a JAX/Equinox model, a loss function, an optional dataset, and metadata tags into a single, serializable object that any L2CO optimizer or rollout can consume. The package ships ready-made task families — analytic black-box benchmarks ([BBOB](https://github.com/bessagroup/bbob-jax) and CEC 2005), quadratic problems, supervised-learning tasks (spiral, MNIST-1D, Gaussian classification), and PINN-style PDE problems (convection, reaction, wave) — together with `create_*_task` factories and `f3dasm` samplers for building experiment datasets.
+
+## Statement of need
+
+Research on learning to optimize and optimizer selection requires evaluating many optimizers across a diverse, well-characterized set of problems — but these problems usually come from incompatible sources with different input ranges, calling conventions, and metadata. `l2co-tasks` standardizes them behind a single `Task` abstraction: inputs are normalized to `[0, 1]^d` and scaled inside the loss, stochasticity and known global minima are tracked explicitly, and every task carries a hashable tag and serializes to a single `.eqx` file. This makes tasks reproducible, portable, and directly pluggable into the [`l2co`](https://github.com/bessagroup/l2co) / [`rl2co`](https://github.com/bessagroup/rl2co) rollout machinery and [`f3dasm`](https://github.com/bessagroup/f3dasm) experiment pipelines.
+
+## Authorship
+
+**Authors**:
+- Martin van der Schelling ([m.p.vanderschelling@tudelft.nl](mailto:m.p.vanderschelling@tudelft.nl))
+
+**Authors affiliation:**
+- Delft University of Technology (Bessa Research Group)
+
+**Maintainer:**
+- Martin van der Schelling ([m.p.vanderschelling@tudelft.nl](mailto:m.p.vanderschelling@tudelft.nl))
+
+**Maintainer affiliation:**
+- Delft University of Technology (Bessa Research Group)
+
+
+## Getting started
+
+`l2co-tasks` is `uv`-managed and depends on an editable install of a sibling [`f3dasm`](https://github.com/bessagroup/f3dasm) checkout, so lay the repositories out side-by-side before syncing:
+
+```bash
+git clone https://github.com/bessagroup/f3dasm.git
+git clone https://github.com/bessagroup/l2co-tasks.git
+cd l2co-tasks
+uv sync
+```
+
+Create a task from one of the factories and evaluate its loss:
+
+```python
+from l2co_tasks import create_bbob_task
+
+task = create_bbob_task(fn_name="sphere", seed=0, dimensionality=2)
+loss = task.loss_fn(task.model)   # model is the [0, 1]^d input vector
+```
+
+Tasks serialize to a single-file `.eqx` format via `Task.save` / `Task.load`. See the [documentation](https://l2co-tasks.readthedocs.io/en/latest/) for the full list of task families and `create_*_task` factories. To build your own task from scratch, see the [Create your own task](./docs/create_task.ipynb) guide.
+
+## Examples and benchmarks
+
+Two runnable notebooks demonstrate and benchmark the core functionality end-to-end. Both run against the installed package after `uv sync` and are rendered in the [documentation](https://l2co-tasks.readthedocs.io/en/latest/):
+
+- [Create your own task](./docs/create_task.ipynb) — builds a `Task` from scratch and exercises the `Task` API (loss evaluation, serialization, metadata).
+- [PINN benchmark tasks](./docs/pinn_tasks.ipynb) — instantiates and evaluates the physics-informed PDE benchmark suite (Helmholtz, Burgers, Euler, Stokes, PK–PD).
+
+## Community Support
+
+If you find any **issues, bugs or problems** with this package, please use the [GitHub issue tracker](https://github.com/bessagroup/l2co-tasks/issues) to report them.
+
+## License
+
+Copyright (c) 2026, Martin van der Schelling
+
+All rights reserved.
+
+This project is licensed under the BSD 3-Clause License. See [LICENSE](https://github.com/bessagroup/l2co-tasks/blob/main/LICENSE) for the full license text.
+
+## Related repositories
+
+This package is part of the L2CO ecosystem developed in the [Bessa Research Group](https://github.com/bessagroup). The repositories below work together:
+
+- [l2co](https://github.com/bessagroup/L2CO) — Learning to Choose Optimizers: a meta-learner that selects an optimizer from problem features before any evaluations, then reassesses that choice from the observed optimization trajectory.
+- [rl2co](https://github.com/bessagroup/rl2co) — Reinforcement Learning to Choose Optimizers: a JAX-based RL agent that dynamically switches between optimizers during a run.
+- [l2co-tasks](https://github.com/bessagroup/l2co-tasks) — Optimization task definitions (BBOB, CEC 2005, PDE, spiral, …) compatible with the L2CO library.
+- [l2co_experiments](https://github.com/bessagroup/l2co_experiments) — Hydra + f3dasm experiment pipelines (dataset creation, training, rollouts, figures) for the L2CO studies.
+- [agentic-l2co](https://github.com/bessagroup/agentic-l2co) — An LLM-agent drop-in replacement for `l2co.L2COModel`, driving two-stage optimizer selection with an Ollama-hosted LLM.
+- [bbob-jax](https://github.com/bessagroup/bbob-jax) — JAX implementations of the BBOB and CEC 2005 black-box optimization benchmark functions.
+- [f3dasm](https://github.com/bessagroup/f3dasm) — Framework for Data-Driven Design and Analysis of Structures and Materials; provides `ExperimentData`, pipelines, and SLURM orchestration.
+

l2co_tasks-0.1.0/pyproject.toml

@@ -0,0 +1,120 @@
+[build-system]
+build-backend = "setuptools.build_meta"
+requires = ["setuptools>=77.0.0", "wheel"]
+
+[project]
+authors = [
+  {email = "m.p.vanderschelling@tudelft.nl", name = "Martin van der Schelling"}
+]
+classifiers = [
+  "Intended Audience :: Science/Research",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "Topic :: Software Development :: Libraries :: Python Modules",
+  "Programming Language :: Python :: 3"
+]
+dependencies = [
+  "bbob-jax>=1.6.4",
+  "cloudpickle>=3.1.1",
+  "equinox>=0.13.0",
+  "f3dasm",
+  "hydra-core>=1.3.2",
+  "jax",
+  "jaxtyping>=0.3.2",
+  "mnist1d",
+  "numpy",
+  "omegaconf",
+  "optax>=0.2.5",
+  "pandas"
+]
+description = "Optimization tasks compatible with the L2CO library"
+keywords = ["optimization", "tasks", "benchmarking"]
+license = "BSD-3-Clause"
+maintainers = [
+  {email = "m.p.vanderschelling@tudelft.nl", name = "Martin van der Schelling"}
+]
+name = "l2co-tasks"
+readme = "README.md"
+requires-python = ">= 3.12"
+version = "0.1.0"
+
+[project.optional-dependencies]
+dev = [
+  "setuptools>=77.0.0",
+  "build",
+  "wheel",
+  "pre-commit>=4.4.0",
+  "ruff>=0.14.10"
+]
+docs = [
+  "mkdocs>=1.6.0",
+  "mkdocs-gen-files>=0.5.0",
+  "mkdocs-include-exclude-files>=0.1.0",
+  "mkdocs-ipynb>=0.1.1",
+  "mkdocs-material>=9.6.20",
+  "mkdocstrings>=0.30.1",
+  "mkdocstrings-python>=1.18.2",
+  "pymdown-extensions>=10.16.1"
+]
+tests = [
+  "pytest>=8.4.1",
+  "pytest-cov>=6.2.1"
+]
+
+[project.urls]
+bugs = "https://github.com/bessagroup/l2co-tasks/issues"
+homepage = "https://github.com/bessagroup/l2co-tasks"
+
+[tool.pytest.ini_options]
+filterwarnings = [
+  "error::DeprecationWarning:l2co_tasks"
+]
+markers = [
+  "slow: test writes a dataset .npz or runs an inner training loop",
+  "requires_f3dasm: test constructs an f3dasm.ExperimentData",
+  "requires_l2co: drives tasks through l2co (skipped unless l2co imports)"
+]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 79
+
+[tool.ruff.lint]
+exclude = [
+  "__pycache__",
+  "dist",
+  "*.egg",
+  "tests",
+  ".git"
+]
+ignore = [
+  "E226",
+  "E3",
+  "E731",
+  "C901",
+  "UP045",
+  "B027",
+  "B012"
+]
+select = [
+  "E", # pycodestyle errors
+  "W", # pycodestyle warnings
+  "F", # Pyflakes
+  "I", # isort
+  "B", # flake8-bugbear
+  "UP" # pyupgrade
+]
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401", "E402"]
+
+[tool.ruff.lint.pydocstyle]
+convention = "numpy"
+
+[tool.setuptools.package-data]
+"l2co_tasks.conf" = ["**/*.yaml"]
+
+[tool.uv]
+package = true
+
+[tool.uv.sources]
+f3dasm = {editable = true, path = "../f3dasm"}

l2co_tasks-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

l2co_tasks-0.1.0/src/l2co_tasks/__init__.py

@@ -0,0 +1,71 @@
+"""
+L2CO Tasks - Optimization tasks compatible with the L2CO library
+"""
+
+#                                                                       Modules
+# =============================================================================
+
+# Local
+from ._src.benchmark_task import (
+    CEC2019Sampler,
+    create_bbob_task,
+    create_cec2005_task,
+)
+from ._src.continue_from_path import retrieve_tasks
+from ._src.euler import EulerTaskSampler, create_euler_task
+from ._src.experimentdata import create_tasks_experimentdata
+from ._src.gaussian_meta import create_gaussian_meta_task
+from ._src.helmholtz import HelmholtzTaskSampler, create_helmholtz_task
+from ._src.inviscid_burgers import (
+    InviscidBurgersTaskSampler,
+    create_inviscid_burgers_task,
+)
+from ._src.mnist1d_task import create_mnist1d_task
+from ._src.pde import PDETaskSampler, create_pde_task
+from ._src.pkpd import PKPDTaskSampler, create_pkpd_task
+from ._src.quadratic_task import create_quadratic_task
+from ._src.spiral_task import create_spiral_task
+from ._src.stokes import StokesTaskSampler, create_stokes_task
+from ._src.task import DatasetDict, Task
+from ._src.task_gaussian_class import create_gaussian_task
+from ._src.viscous_burgers import (
+    ViscousBurgersTaskSampler,
+    create_viscous_burgers_task,
+)
+
+#                                                        Authorship and Credits
+# =============================================================================
+__author__ = "Martin van der Schelling (m.p.vanderschelling@tudelft.nl)"
+__credits__ = ["Martin van der Schelling"]
+__status__ = "Stable"
+#
+# =============================================================================
+
+__all__ = [
+    "CEC2019Sampler",
+    "DatasetDict",
+    "EulerTaskSampler",
+    "HelmholtzTaskSampler",
+    "InviscidBurgersTaskSampler",
+    "PDETaskSampler",
+    "PKPDTaskSampler",
+    "StokesTaskSampler",
+    "Task",
+    "ViscousBurgersTaskSampler",
+    "create_bbob_task",
+    "create_cec2005_task",
+    "create_euler_task",
+    "create_gaussian_meta_task",
+    "create_gaussian_task",
+    "create_helmholtz_task",
+    "create_inviscid_burgers_task",
+    "create_mnist1d_task",
+    "create_pde_task",
+    "create_pkpd_task",
+    "create_quadratic_task",
+    "create_spiral_task",
+    "create_stokes_task",
+    "create_tasks_experimentdata",
+    "create_viscous_burgers_task",
+    "retrieve_tasks",
+]

l2co_tasks-0.1.0/src/l2co_tasks/_src/__init__.py

@@ -0,0 +1,5 @@
+"""Internal implementation of ``l2co_tasks``.
+
+Public API is re-exported from :mod:`l2co_tasks`; do not import from this
+subpackage directly in external code.
+"""

l2co_tasks-0.1.0/src/l2co_tasks/_src/_io.py

@@ -0,0 +1,56 @@
+"""Dataset loading and saving helpers."""
+
+#                                                                       Modules
+# =============================================================================
+
+# Standard
+from pathlib import Path
+
+# Third-party
+import jax.numpy as jnp
+import numpy as np
+
+#                                                          Authorship & Credits
+# =============================================================================
+__author__ = "Martin van der Schelling (M.P.vanderSchelling@tudelft.nl)"
+__credits__ = ["Martin van der Schelling"]
+__status__ = "Stable"
+# =============================================================================
+
+
+def load_dataset(path: str | Path | None) -> dict[str, jnp.ndarray]:
+    """Load dataset from npz file.
+
+    Parameters
+    ----------
+    path : str | Path | None
+        Path to the dataset file. If None, returns empty dict.
+
+    Returns
+    -------
+    dict[str, jnp.ndarray]
+        Loaded dataset dictionary.
+    """
+    if path is None:
+        return {}
+    with np.load(Path(path).with_suffix(".npz")) as loaded_data:
+        return {k: jnp.asarray(loaded_data[k]) for k in loaded_data.files}
+
+
+def save_dataset(dataset: dict[str, jnp.ndarray], path: str | Path) -> None:
+    """Save a dataset of arrays to an ``.npz`` file.
+
+    Creates the parent directory if needed. The write-side twin of
+    :func:`load_dataset`; the task factories delegate here rather than
+    each carrying its own copy.
+
+    Parameters
+    ----------
+    dataset : dict[str, jnp.ndarray]
+        Dictionary of arrays to persist.
+    path : str | Path
+        Destination file path.
+    """
+    _path = Path(path)
+    _path.parent.mkdir(parents=True, exist_ok=True)
+    jnp.savez(path, **dataset)