From 2bb1a5dc1139daaa23f0ec9d4f20c65b1ffca1ae Mon Sep 17 00:00:00 2001 From: Fabrice Normandin Date: Mon, 26 Aug 2024 08:33:58 -0400 Subject: [PATCH] Automatically create schemas for yaml config used by Hydra (#7) * WIP: Automatically create yaml schemas for configs Signed-off-by: Fabrice Normandin * Working (mostly)! Signed-off-by: Fabrice Normandin * Add a .gitignore entry for schemas Signed-off-by: Fabrice Normandin * Associate files to schemas via vscode settings Signed-off-by: Fabrice Normandin * Improve descriptions Signed-off-by: Fabrice Normandin * Update notes / todos Signed-off-by: Fabrice Normandin * Fix doctest Signed-off-by: Fabrice Normandin * Save progress so far (lots more to do :() Signed-off-by: Fabrice Normandin * Ignore most hydra errors for now Signed-off-by: Fabrice Normandin * Fix issues with logging verbosity Signed-off-by: Fabrice Normandin * Rename functions and change signature Signed-off-by: Fabrice Normandin * Make it possible to change only one file Signed-off-by: Fabrice Normandin * Add "format on save" option to devcontainer Signed-off-by: Fabrice Normandin * Fix bug in experiment/example.yaml Signed-off-by: Fabrice Normandin * Format a yaml file Signed-off-by: Fabrice Normandin * Regenerate schema if config changes Signed-off-by: Fabrice Normandin * Remove unused imports in example.py Signed-off-by: Fabrice Normandin * Add nice schema for `defaults` list! Signed-off-by: Fabrice Normandin * Fix issue with unit test for auto schema Signed-off-by: Fabrice Normandin * Add a 'features/auto_schema.md' file with a demo Signed-off-by: Fabrice Normandin * Create schemas inside `main`, tweak defaults Signed-off-by: Fabrice Normandin * Handle sub-entries with _target_ Signed-off-by: Fabrice Normandin * Fix the errors in the test configs for auto_shema Signed-off-by: Fabrice Normandin * Remove absolute path in regression files Signed-off-by: Fabrice Normandin * Add tests to increase code coverage Signed-off-by: Fabrice Normandin * Fix issue when writing yaml schema header Signed-off-by: Fabrice Normandin * Add test to cover changes in `main.py` Signed-off-by: Fabrice Normandin * "fix" type-checking errors in auto_schema.py Signed-off-by: Fabrice Normandin --------- Signed-off-by: Fabrice Normandin --- .devcontainer/devcontainer.json | 7 +- .gitignore | 1 + docs/SUMMARY.md | 2 + docs/features/auto_schema.md | 18 + mkdocs.yml | 58 +- project/__init__.py | 2 +- project/algorithms/callbacks/callback.py | 2 +- .../callbacks/classification_metrics.py | 2 +- .../callbacks/samples_per_second.py | 2 +- project/algorithms/example.py | 3 +- project/algorithms/jax_example.py | 2 +- project/algorithms/no_op.py | 2 +- project/algorithms/testsuites/algorithm.py | 4 +- .../algorithms/testsuites/algorithm_tests.py | 4 +- project/configs/config.yaml | 1 - project/configs/datamodule/cifar10.yaml | 3 +- project/configs/datamodule/fashion_mnist.yaml | 3 +- project/configs/datamodule/imagenet.yaml | 1 + project/configs/datamodule/imagenet32.yaml | 3 +- project/configs/datamodule/inaturalist.yaml | 3 +- project/configs/datamodule/mnist.yaml | 3 +- .../experiment/cluster_sweep_example.yaml | 13 +- project/configs/experiment/example.yaml | 6 +- project/configs/trainer/default.yaml | 1 - project/conftest.py | 4 +- project/datamodules/datamodules_test.py | 2 +- .../image_classification/cifar10.py | 2 +- .../image_classification.py | 4 +- .../image_classification/imagenet.py | 4 +- .../image_classification/imagenet32.py | 2 +- .../image_classification/inaturalist.py | 2 +- .../datamodules/image_classification/mnist.py | 2 +- project/datamodules/vision.py | 4 +- project/experiment.py | 4 +- project/main.py | 19 +- project/main_test.py | 12 + project/networks/layers/layers.py | 2 +- project/networks/layers/sequential.py | 2 +- project/utils/auto_schema.py | 1093 +++++++++++++++++ project/utils/auto_schema_test.py | 94 ++ project/utils/auto_schema_test/nested.json | 129 ++ project/utils/auto_schema_test/nested.yaml | 6 + project/utils/auto_schema_test/partial.json | 99 ++ project/utils/auto_schema_test/partial.yaml | 3 + .../utils/auto_schema_test/with_defaults.json | 101 ++ .../utils/auto_schema_test/with_defaults.yaml | 4 + .../auto_schema_test/with_package_global.json | 101 ++ .../auto_schema_test/with_package_global.yaml | 4 + .../utils/auto_schema_test/with_target.json | 101 ++ .../utils/auto_schema_test/with_target.yaml | 3 + project/utils/hydra_config_utils.py | 4 +- .../utils/{types => typing_utils}/__init__.py | 0 .../{types => typing_utils}/protocols.py | 0 project/utils/utils.py | 2 +- pyproject.toml | 10 +- requirements-dev.lock | 326 ++--- requirements.lock | 276 +++-- 57 files changed, 2196 insertions(+), 371 deletions(-) create mode 100644 docs/features/auto_schema.md create mode 100644 project/utils/auto_schema.py create mode 100644 project/utils/auto_schema_test.py create mode 100644 project/utils/auto_schema_test/nested.json create mode 100644 project/utils/auto_schema_test/nested.yaml create mode 100644 project/utils/auto_schema_test/partial.json create mode 100644 project/utils/auto_schema_test/partial.yaml create mode 100644 project/utils/auto_schema_test/with_defaults.json create mode 100644 project/utils/auto_schema_test/with_defaults.yaml create mode 100644 project/utils/auto_schema_test/with_package_global.json create mode 100644 project/utils/auto_schema_test/with_package_global.yaml create mode 100644 project/utils/auto_schema_test/with_target.json create mode 100644 project/utils/auto_schema_test/with_target.yaml rename project/utils/{types => typing_utils}/__init__.py (100%) rename project/utils/{types => typing_utils}/protocols.py (100%) diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index d49eb4ec..8828209d 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -30,6 +30,10 @@ "editor.defaultFormatter": "charliermarsh.ruff", "editor.formatOnType": true }, + "editor.rulers": [ + 99 + ], + "editor.formatOnSave": true, "files.exclude": { "**/.git": true, "**/.svn": true, @@ -61,7 +65,8 @@ "GitHub.copilot", "knowsuchagency.pdm-task-provider", "GitHub.copilot-chat", - "mutantdino.resourcemonitor" + "mutantdino.resourcemonitor", + "Gruntfuggly.todo-tree" ] } }, diff --git a/.gitignore b/.gitignore index 34e527c2..da972e21 100644 --- a/.gitignore +++ b/.gitignore @@ -169,3 +169,4 @@ lightning_logs **.npz # .python-version .testmondata* +.schemas diff --git a/docs/SUMMARY.md b/docs/SUMMARY.md index 716ea5fd..17217684 100644 --- a/docs/SUMMARY.md +++ b/docs/SUMMARY.md @@ -3,6 +3,8 @@ * overview/*.md * Getting Started * getting_started/*.md + * Features + * features/*.md * Reference * reference/* * Examples diff --git a/docs/features/auto_schema.md b/docs/features/auto_schema.md new file mode 100644 index 00000000..821ad11e --- /dev/null +++ b/docs/features/auto_schema.md @@ -0,0 +1,18 @@ +# Auto Schema for Hydra Configs + +This project template comes with a really neat feature: Your [Hydra](https://hydra.cc) config files automatically get a [Schema](https://json-schema.org/) associated with them. + +This greatly improves the experience of developing a project with Hydra: + +- Saves you time by preventing errors caused by unexpected keys in your config files, or values that are of the wrong type + This can often happen after moving files or renaming a function, for example. +- While writing a config file you get to see: + - the list of available configuration options in a given config + - the default values for each value + - the documentation for each value (taken from the source code of the function!) + +Here's a quick demo of what this looks like in practice: + +![type:video](https://github.com/user-attachments/assets/08f52d47-ebba-456d-95ef-ac9525d8e983) + +Here we have a config that will be used to configure the `lightning.Trainer` class, but any config file in the project will also get a schema automatically, even if it doesn't have a `"_target_"` key directly! diff --git a/mkdocs.yml b/mkdocs.yml index 346c9b0b..2ee612e1 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -14,32 +14,38 @@ markdown_extensions: - pymdownx.superfences plugins: -- search -- literate-nav: - nav_file: SUMMARY.md -- awesome-pages -- gen-files: - # https://oprypin.github.io/mkdocs-gen-files/#usage - scripts: - - docs/generate_reference_docs.py -- mkdocstrings: - handlers: - python: - import: - - https://docs.python-requests.org/en/master/objects.inv - options: - docstring_style: google - members_order: source - annotations_path: brief - show_docstring_attributes: true - modernize_annotations: true - show_source: false - show_submodules: false - separate_signature: true - signature_crossrefs: true - show_signature_annotations: true - allow_inspection: true - + - search + - literate-nav: + nav_file: SUMMARY.md + - awesome-pages + - gen-files: + # https://oprypin.github.io/mkdocs-gen-files/#usage + scripts: + - docs/generate_reference_docs.py + - mkdocstrings: + handlers: + python: + import: + - https://docs.python-requests.org/en/master/objects.inv + - https://lightning.ai/docs/pytorch/stable/objects.inv + options: + docstring_style: google + members_order: source + annotations_path: brief + show_docstring_attributes: true + modernize_annotations: true + show_source: false + show_submodules: false + separate_signature: true + signature_crossrefs: true + show_signature_annotations: true + allow_inspection: true + - mkdocs-video: + is_video: True + video_muted: True + video_controls: True + css_style: + width: "100%" # todo: take a look at https://github.com/drivendataorg/cookiecutter-data-science/blob/master/docs/mkdocs.yml # - admonition # - pymdownx.details diff --git a/project/__init__.py b/project/__init__.py index 8484d37f..cb90d760 100644 --- a/project/__init__.py +++ b/project/__init__.py @@ -4,7 +4,7 @@ from .utils.hydra_utils import patched_safe_name # noqa # from .networks import FcNet -from .utils.types import DataModule +from .utils.typing_utils import DataModule add_configs_to_hydra_store() diff --git a/project/algorithms/callbacks/callback.py b/project/algorithms/callbacks/callback.py index 3f6c7bfd..f8635ff6 100644 --- a/project/algorithms/callbacks/callback.py +++ b/project/algorithms/callbacks/callback.py @@ -10,7 +10,7 @@ from lightning import pytorch as pl from typing_extensions import TypeVar, override -from project.utils.types import NestedMapping +from project.utils.typing_utils import NestedMapping from project.utils.utils import get_log_dir logger = get_logger(__name__) diff --git a/project/algorithms/callbacks/classification_metrics.py b/project/algorithms/callbacks/classification_metrics.py index 95fbf149..8988216d 100644 --- a/project/algorithms/callbacks/classification_metrics.py +++ b/project/algorithms/callbacks/classification_metrics.py @@ -10,7 +10,7 @@ from typing_extensions import NotRequired, Required, override from project.algorithms.callbacks.callback import BatchType, Callback -from project.utils.types.protocols import ClassificationDataModule +from project.utils.typing_utils.protocols import ClassificationDataModule logger = get_logger(__name__) diff --git a/project/algorithms/callbacks/samples_per_second.py b/project/algorithms/callbacks/samples_per_second.py index c7399063..853c6b06 100644 --- a/project/algorithms/callbacks/samples_per_second.py +++ b/project/algorithms/callbacks/samples_per_second.py @@ -7,7 +7,7 @@ from typing_extensions import override from project.algorithms.callbacks.callback import BatchType, Callback, StepOutputType -from project.utils.types import is_sequence_of +from project.utils.typing_utils import is_sequence_of class MeasureSamplesPerSecondCallback(Callback[BatchType, StepOutputType]): diff --git a/project/algorithms/example.py b/project/algorithms/example.py index 015f88d4..7ead41c2 100644 --- a/project/algorithms/example.py +++ b/project/algorithms/example.py @@ -9,12 +9,11 @@ from typing import Any, Literal import torch -from hydra_zen.typing import HydraPartialBuilds, Partial, PartialBuilds, ZenPartialBuilds # noqa from lightning import LightningModule from omegaconf import DictConfig from torch import Tensor from torch.nn import functional as F -from torch.optim import Optimizer +from torch.optim.optimizer import Optimizer from project.configs.algorithm.optimizer import AdamConfig from project.datamodules.image_classification import ImageClassificationDataModule diff --git a/project/algorithms/jax_example.py b/project/algorithms/jax_example.py index 58bfb28e..947dfd6e 100644 --- a/project/algorithms/jax_example.py +++ b/project/algorithms/jax_example.py @@ -19,7 +19,7 @@ ImageClassificationDataModule, ) from project.datamodules.image_classification.mnist import MNISTDataModule -from project.utils.types.protocols import ClassificationDataModule +from project.utils.typing_utils.protocols import ClassificationDataModule os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" diff --git a/project/algorithms/no_op.py b/project/algorithms/no_op.py index 11bfb898..3824448d 100644 --- a/project/algorithms/no_op.py +++ b/project/algorithms/no_op.py @@ -5,7 +5,7 @@ from torch import nn from project.algorithms.callbacks.samples_per_second import MeasureSamplesPerSecondCallback -from project.utils.types.protocols import DataModule +from project.utils.typing_utils.protocols import DataModule class NoOp(LightningModule): diff --git a/project/algorithms/testsuites/algorithm.py b/project/algorithms/testsuites/algorithm.py index 8b65c91b..e2e84a65 100644 --- a/project/algorithms/testsuites/algorithm.py +++ b/project/algorithms/testsuites/algorithm.py @@ -5,8 +5,8 @@ from torch import Tensor from typing_extensions import NotRequired, TypeVar -from project.utils.types import PyTree -from project.utils.types.protocols import DataModule, Module +from project.utils.typing_utils import PyTree +from project.utils.typing_utils.protocols import DataModule, Module class StepOutputDict(TypedDict, total=False): diff --git a/project/algorithms/testsuites/algorithm_tests.py b/project/algorithms/testsuites/algorithm_tests.py index 60228948..fcdd70f7 100644 --- a/project/algorithms/testsuites/algorithm_tests.py +++ b/project/algorithms/testsuites/algorithm_tests.py @@ -19,8 +19,8 @@ ParametrizedFixture, seeded_rng, ) -from project.utils.types import PyTree, is_sequence_of -from project.utils.types.protocols import DataModule +from project.utils.typing_utils import PyTree, is_sequence_of +from project.utils.typing_utils.protocols import DataModule logger = get_logger(__name__) diff --git a/project/configs/config.yaml b/project/configs/config.yaml index 3fa14cf3..19150416 100644 --- a/project/configs/config.yaml +++ b/project/configs/config.yaml @@ -15,5 +15,4 @@ defaults: # experiment configs allow for version control of specific hyperparameters # e.g. best hyperparameters for given model and datamodule - experiment: null - # name: "${hydra:runtime.choices.algorithm}-${hydra:runtime.choices.network}-${hydra:runtime.choices.datamodule}" diff --git a/project/configs/datamodule/cifar10.yaml b/project/configs/datamodule/cifar10.yaml index 0ca045ea..2410ef7c 100644 --- a/project/configs/datamodule/cifar10.yaml +++ b/project/configs/datamodule/cifar10.yaml @@ -1,5 +1,6 @@ defaults: -- vision + - vision + - _self_ _target_: project.datamodules.CIFAR10DataModule data_dir: ${constant:torchvision_dir,DATA_DIR} batch_size: 128 diff --git a/project/configs/datamodule/fashion_mnist.yaml b/project/configs/datamodule/fashion_mnist.yaml index f2038d2d..472a4d96 100644 --- a/project/configs/datamodule/fashion_mnist.yaml +++ b/project/configs/datamodule/fashion_mnist.yaml @@ -1,3 +1,4 @@ defaults: -- mnist + - mnist + - _self_ _target_: project.datamodules.FashionMNISTDataModule diff --git a/project/configs/datamodule/imagenet.yaml b/project/configs/datamodule/imagenet.yaml index 3e82c78b..23804087 100644 --- a/project/configs/datamodule/imagenet.yaml +++ b/project/configs/datamodule/imagenet.yaml @@ -1,4 +1,5 @@ defaults: - vision + - _self_ _target_: project.datamodules.ImageNetDataModule # todo: add good configuration options here. diff --git a/project/configs/datamodule/imagenet32.yaml b/project/configs/datamodule/imagenet32.yaml index 208119be..ace5c8ae 100644 --- a/project/configs/datamodule/imagenet32.yaml +++ b/project/configs/datamodule/imagenet32.yaml @@ -1,5 +1,6 @@ defaults: -- vision + - vision + - _self_ _target_: project.datamodules.ImageNet32DataModule data_dir: ${constant:SCRATCH} val_split: -1 diff --git a/project/configs/datamodule/inaturalist.yaml b/project/configs/datamodule/inaturalist.yaml index 5670be6b..d3621b0f 100644 --- a/project/configs/datamodule/inaturalist.yaml +++ b/project/configs/datamodule/inaturalist.yaml @@ -1,5 +1,6 @@ defaults: -- vision + - vision + - _self_ _target_: project.datamodules.INaturalistDataModule version: "2021_train" target_type: "full" diff --git a/project/configs/datamodule/mnist.yaml b/project/configs/datamodule/mnist.yaml index a7554ec3..625b1ad9 100644 --- a/project/configs/datamodule/mnist.yaml +++ b/project/configs/datamodule/mnist.yaml @@ -1,5 +1,6 @@ defaults: -- vision + - vision + - _self_ _target_: project.datamodules.MNISTDataModule data_dir: ${constant:torchvision_dir,DATA_DIR} normalize: True diff --git a/project/configs/experiment/cluster_sweep_example.yaml b/project/configs/experiment/cluster_sweep_example.yaml index 03510ac8..96992b8d 100644 --- a/project/configs/experiment/cluster_sweep_example.yaml +++ b/project/configs/experiment/cluster_sweep_example.yaml @@ -38,13 +38,13 @@ hydra: subdir: ${hydra.job.id}/task${oc.env:SLURM_PROCID} launcher: - array_parallelism: 4 # max num of jobs to run in parallel + array_parallelism: 4 # max num of jobs to run in parallel # note: Can't currently use this argument with the submitit launcher plugin. We would need to # create a subclass to add support for it. # ntasks_per_gpu: 2 # Other things to pass to `sbatch`: additional_parameters: - time: 0-00:10:00 # maximum wall time allocated for the job (D-HH:MM:SS) + time: 0-00:10:00 # maximum wall time allocated for the job (D-HH:MM:SS) # requeue: null # requeue job if it fails ## A list of commands to add to the generated sbatch script before running srun: # setup: @@ -61,12 +61,11 @@ hydra: name: "${name}" version: 1 - algorithm: - # BUG: Getting a weird bug with TPE: KeyError in `dum_below_trials = [...]` at line 397. + # BUG: Getting a weird bug with TPE: KeyError in `dum_below_trials = [...]` at line 397. type: random config: - seed: 1 + seed: 1 worker: n_workers: ${hydra.launcher.array_parallelism} @@ -77,5 +76,5 @@ hydra: type: legacy use_hydra_path: false database: - type: pickleddb - host: "logs/${name}/multiruns/database.pkl" + type: pickleddb + host: "logs/${name}/multiruns/database.pkl" diff --git a/project/configs/experiment/example.yaml b/project/configs/experiment/example.yaml index e11f37c3..73779651 100644 --- a/project/configs/experiment/example.yaml +++ b/project/configs/experiment/example.yaml @@ -18,14 +18,12 @@ name: example seed: ${oc.env:SLURM_PROCID,12345} algorithm: - hp: - optimizer: - lr: 0.002 + optimizer_config: + lr: 0.002 datamodule: batch_size: 64 - trainer: min_epochs: 1 max_epochs: 10 diff --git a/project/configs/trainer/default.yaml b/project/configs/trainer/default.yaml index f4cec3bc..de5d21be 100644 --- a/project/configs/trainer/default.yaml +++ b/project/configs/trainer/default.yaml @@ -1,5 +1,4 @@ _target_: lightning.Trainer - logger: null accelerator: auto strategy: auto diff --git a/project/conftest.py b/project/conftest.py index a9218705..88ea7d4d 100644 --- a/project/conftest.py +++ b/project/conftest.py @@ -49,8 +49,8 @@ default_marks_for_config_name, seeded_rng, ) -from project.utils.types import is_sequence_of -from project.utils.types.protocols import ( +from project.utils.typing_utils import is_sequence_of +from project.utils.typing_utils.protocols import ( DataModule, ) diff --git a/project/datamodules/datamodules_test.py b/project/datamodules/datamodules_test.py index a2f07815..cd733d3d 100644 --- a/project/datamodules/datamodules_test.py +++ b/project/datamodules/datamodules_test.py @@ -18,7 +18,7 @@ ) from project.datamodules.vision import VisionDataModule from project.utils.testutils import run_for_all_datamodules -from project.utils.types import is_sequence_of +from project.utils.typing_utils import is_sequence_of # @use_overrides(["datamodule.num_workers=0"]) diff --git a/project/datamodules/image_classification/cifar10.py b/project/datamodules/image_classification/cifar10.py index 91cbc92e..99f41c5f 100644 --- a/project/datamodules/image_classification/cifar10.py +++ b/project/datamodules/image_classification/cifar10.py @@ -10,7 +10,7 @@ ImageClassificationDataModule, ) from project.datamodules.vision import VisionDataModule -from project.utils.types import C, H, W +from project.utils.typing_utils import C, H, W def cifar10_train_transforms(): diff --git a/project/datamodules/image_classification/image_classification.py b/project/datamodules/image_classification/image_classification.py index cdb1ad58..18fed50d 100644 --- a/project/datamodules/image_classification/image_classification.py +++ b/project/datamodules/image_classification/image_classification.py @@ -6,8 +6,8 @@ from torchvision.tv_tensors import Image from project.datamodules.vision import VisionDataModule -from project.utils.types import C, H, W -from project.utils.types.protocols import ClassificationDataModule +from project.utils.typing_utils import C, H, W +from project.utils.typing_utils.protocols import ClassificationDataModule # todo: need to decide whether this should be a base class or just a protocol. # - IF this is a protocol, then we can't use issubclass with it, so it can't be used in the diff --git a/project/datamodules/image_classification/imagenet.py b/project/datamodules/image_classification/imagenet.py index 158e6669..b32debec 100644 --- a/project/datamodules/image_classification/imagenet.py +++ b/project/datamodules/image_classification/imagenet.py @@ -23,8 +23,8 @@ from project.datamodules.vision import VisionDataModule from project.utils.env_vars import DATA_DIR, NETWORK_DIR, NUM_WORKERS -from project.utils.types import C, H, W -from project.utils.types.protocols import Module +from project.utils.typing_utils import C, H, W +from project.utils.typing_utils.protocols import Module logger = get_logger(__name__) diff --git a/project/datamodules/image_classification/imagenet32.py b/project/datamodules/image_classification/imagenet32.py index 874a053a..698bdb9b 100644 --- a/project/datamodules/image_classification/imagenet32.py +++ b/project/datamodules/image_classification/imagenet32.py @@ -19,7 +19,7 @@ from project.datamodules.vision import VisionDataModule from project.utils.env_vars import DATA_DIR, SCRATCH -from project.utils.types import C, H, W +from project.utils.typing_utils import C, H, W logger = getLogger(__name__) diff --git a/project/datamodules/image_classification/inaturalist.py b/project/datamodules/image_classification/inaturalist.py index ca67f7ed..58eb8d52 100644 --- a/project/datamodules/image_classification/inaturalist.py +++ b/project/datamodules/image_classification/inaturalist.py @@ -13,7 +13,7 @@ ImageClassificationDataModule, ) from project.utils.env_vars import DATA_DIR, NUM_WORKERS, SLURM_TMPDIR -from project.utils.types import C, H, W +from project.utils.typing_utils import C, H, W logger = get_logger(__name__) diff --git a/project/datamodules/image_classification/mnist.py b/project/datamodules/image_classification/mnist.py index 1b96f850..fecf1753 100644 --- a/project/datamodules/image_classification/mnist.py +++ b/project/datamodules/image_classification/mnist.py @@ -11,7 +11,7 @@ from project.datamodules.image_classification.image_classification import ( ImageClassificationDataModule, ) -from project.utils.types import C, H, W +from project.utils.typing_utils import C, H, W def mnist_train_transforms(): diff --git a/project/datamodules/vision.py b/project/datamodules/vision.py index 0d06d4ec..50d8dd12 100644 --- a/project/datamodules/vision.py +++ b/project/datamodules/vision.py @@ -17,8 +17,8 @@ from torchvision.tv_tensors import Image, set_return_type from project.utils.env_vars import DATA_DIR, NUM_WORKERS -from project.utils.types import C, H, W -from project.utils.types.protocols import DataModule +from project.utils.typing_utils import C, H, W +from project.utils.typing_utils.protocols import DataModule logger = get_logger(__name__) diff --git a/project/experiment.py b/project/experiment.py index 081788a3..4eea56ab 100644 --- a/project/experiment.py +++ b/project/experiment.py @@ -23,8 +23,8 @@ ImageClassificationDataModule, ) from project.utils.hydra_utils import get_outer_class -from project.utils.types import Dataclass -from project.utils.types.protocols import DataModule, Module +from project.utils.typing_utils import Dataclass +from project.utils.typing_utils.protocols import DataModule, Module from project.utils.utils import validate_datamodule logger = get_logger(__name__) diff --git a/project/main.py b/project/main.py index 14871b94..f2ea9585 100644 --- a/project/main.py +++ b/project/main.py @@ -17,6 +17,7 @@ from project.configs.config import Config from project.experiment import Experiment, setup_experiment +from project.utils.env_vars import REPO_ROOTDIR from project.utils.hydra_utils import resolve_dictconfig from project.utils.utils import print_config @@ -25,17 +26,31 @@ os.environ["CUDA_VISIBLE_DEVICES"] = "0" logger = get_logger(__name__) -project_name = Path(__file__).parent.name +PROJECT_NAME = Path(__file__).parent.name @hydra.main( - config_path="pkg://project.configs", + config_path=f"pkg://{PROJECT_NAME}.configs", config_name="config", version_base="1.2", ) def main(dict_config: DictConfig) -> dict: """Main entry point for training a model.""" print_config(dict_config, resolve=False) + + from project.utils.auto_schema import add_schemas_to_all_hydra_configs + + # Note: running this should take ~5 seconds the first time and <1s after that. + add_schemas_to_all_hydra_configs( + config_files=None, + repo_root=REPO_ROOTDIR, + configs_dir=REPO_ROOTDIR / PROJECT_NAME / "configs", + regen_schemas=False, + stop_on_error=False, + quiet=True, + add_headers=False, + ) + config: Config = resolve_dictconfig(dict_config) experiment: Experiment = setup_experiment(config) diff --git a/project/main_test.py b/project/main_test.py index fdff471b..8a688b72 100644 --- a/project/main_test.py +++ b/project/main_test.py @@ -4,12 +4,15 @@ import shutil import hydra_zen +from omegaconf import DictConfig from project.algorithms.example import ExampleAlgorithm from project.configs.config import Config from project.conftest import use_overrides from project.datamodules.image_classification.cifar10 import CIFAR10DataModule +from .main import main + def test_jax_can_use_the_GPU(): import jax.numpy @@ -38,6 +41,15 @@ def test_defaults(experiment_config: Config) -> None: ) +@use_overrides(["seed=1 +trainer.fast_dev_run=True"]) +def test_fast_dev_run(experiment_dictconfig: DictConfig): + result = main(experiment_dictconfig) + assert isinstance(result, dict) + assert result["type"] == "objective" + assert isinstance(result["name"], str) + assert isinstance(result["value"], float) + + # TODO: Add some more integration tests: # - running sweeps from Hydra! # - using the slurm launcher! diff --git a/project/networks/layers/layers.py b/project/networks/layers/layers.py index d6496cb4..a82fd1a0 100644 --- a/project/networks/layers/layers.py +++ b/project/networks/layers/layers.py @@ -13,7 +13,7 @@ from torch import Tensor, nn from typing_extensions import ParamSpec -from project.utils.types import Module, OutT, T, is_sequence_of +from project.utils.typing_utils import Module, OutT, T, is_sequence_of P = ParamSpec("P", default=[Tensor]) diff --git a/project/networks/layers/sequential.py b/project/networks/layers/sequential.py index 091adb44..b4427257 100644 --- a/project/networks/layers/sequential.py +++ b/project/networks/layers/sequential.py @@ -11,7 +11,7 @@ from torch._jit_internal import _copy_to_script_wrapper from typing_extensions import TypeVar -from project.utils.types import Module +from project.utils.typing_utils import Module ModuleType = TypeVar("ModuleType", bound=Module[..., Any], default=Module[[Tensor], Tensor]) diff --git a/project/utils/auto_schema.py b/project/utils/auto_schema.py new file mode 100644 index 00000000..68daa6e2 --- /dev/null +++ b/project/utils/auto_schema.py @@ -0,0 +1,1093 @@ +"""Scripts that creates Schemas for hydra config files. + +This is very helpful when using Hydra! It shows the user what options are available, along with their +description and default values, and displays errors if you have config files with invalid values. + +## todos +- [ ] Add schemas for all the nested dict entries in a config file if they have a _target_ (currently just the first level). +- [ ] Modify the schema to support omegaconf directives like ${oc.env:VAR_NAME} and our custom directives like ${instance_attr} and so on. +- [ ] todo: Make a hydra plugin that creates the schemas for configs when hydra is loading stuff. +""" + +from __future__ import annotations + +import argparse +import copy +import dataclasses +import datetime +import importlib +import inspect +import json +import logging +import os.path +import shutil +import subprocess +import sys +import typing +import warnings +from collections.abc import Callable, Mapping, MutableMapping, Sequence +from logging import getLogger as get_logger +from pathlib import Path +from typing import Any, Literal, TypedDict, TypeVar + +import docstring_parser as dp +import flax +import flax.linen +import flax.struct +import hydra.errors +import hydra.utils +import hydra_zen +import lightning.pytorch.callbacks +import omegaconf +import pydantic +import pydantic.schema +import rich.logging +import tqdm +from hydra._internal.config_loader_impl import ConfigLoaderImpl +from hydra._internal.utils import create_config_search_path +from hydra.core.config_search_path import ConfigSearchPath +from hydra.plugins.search_path_plugin import SearchPathPlugin +from hydra.types import RunMode +from omegaconf import DictConfig, OmegaConf +from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue +from pydantic_core import core_schema +from tqdm.rich import tqdm_rich +from typing_extensions import NotRequired, Required + +from project.utils.env_vars import REPO_ROOTDIR + +logger = get_logger(__name__) + +_SpecialEntries = TypedDict( + "_SpecialEntries", {"$defs": dict[str, dict], "$schema": str}, total=False +) + + +class PropertySchema(_SpecialEntries, total=False): + title: str + type: str + description: str + default: Any + examples: list[str] + deprecated: bool + readOnly: bool + writeOnly: bool + const: Any + enum: list[Any] + + +# S = TypeVar("S", bound=PropertySchema, covariant=True) + + +class OneOf(TypedDict): + oneOf: Sequence[PropertySchema | ArrayPropertySchema | ObjectSchema | StringPropertySchema] + + +class ArrayPropertySchema(PropertySchema, total=False): + type: Literal["array"] + items: Required[PropertySchema | OneOf] + minItems: int + maxItems: int + uniqueItems: bool + + +class StringPropertySchema(PropertySchema, total=False): + type: Literal["string"] + pattern: str + + +class _PropertyNames(TypedDict): + pattern: str + + +class ObjectSchema(PropertySchema, total=False): + type: Literal["object"] + # annoying that we have to include the subclasses here! + properties: MutableMapping[ + str, PropertySchema | ArrayPropertySchema | StringPropertySchema | ObjectSchema + ] + patternProperties: Mapping[str, PropertySchema | StringPropertySchema] + propertyNames: _PropertyNames + minProperties: int + maxProperties: int + + +class Schema(TypedDict, total=False): + # "$defs": + title: str + description: str + type: str + + # annoying that we have to include the subclasses here! + properties: Required[ + MutableMapping[ + str, PropertySchema | ArrayPropertySchema | StringPropertySchema | ObjectSchema + ] + ] + required: NotRequired[list[str]] + + additionalProperties: NotRequired[bool] + + dependentRequired: NotRequired[MutableMapping[str, list[str]]] + """ https://json-schema.org/understanding-json-schema/reference/conditionals#dependentRequired """ + + +HYDRA_CONFIG_SCHEMA = Schema( + title="Default Schema for any Hydra config file.", + description="Schema created by the `auto_schema.py` script.", + properties={ + "defaults": ArrayPropertySchema( + title="Hydra defaults", + description="Hydra defaults for this config. See https://hydra.cc/docs/advanced/defaults_list/", + type="array", + items=OneOf( + oneOf=[ + ObjectSchema( + type="object", + propertyNames={"pattern": r"^(override\s*)?(/?\w*)+$"}, + patternProperties={ + # todo: support package directives with @? + # todo: Create a enum schema using the available options in the config group! + # override network: something -> `something` value should be in the available options for network. + r"^(override\s*)?(/?\w*)*$": StringPropertySchema( + type="string", pattern=r"\w*(.yaml|.yml)?$" + ), + }, + minProperties=1, + maxProperties=1, + ), + StringPropertySchema(type="string", pattern=r"^\w+(.yaml|.yml)?$"), + ObjectSchema( + type="object", + propertyNames=_PropertyNames(pattern=r"^(override\s*)?(/?\w*)+$"), + patternProperties={ + r"^(override\s*)?(/?\w*)*$": PropertySchema(type="null"), + }, + minProperties=1, + maxProperties=1, + ), + ], + ), + uniqueItems=True, + ), + "_target_": StringPropertySchema( + type="string", + title="Target", + description="Target to instantiate.\nSee https://hydra.cc/docs/advanced/instantiate_objects/overview/", + ), + "_convert_": StringPropertySchema( + type="string", + enum=["none", "partial", "object", "all"], + title="Convert", + description="See https://hydra.cc/docs/advanced/instantiate_objects/overview/#parameter-conversion-strategies", + ), + "_partial_": PropertySchema( + type="boolean", + title="Partial", + description=( + "Whether this config calls the target function when instantiated, or creates " + "a `functools.partial` that will call the target.\n" + "See: https://hydra.cc/docs/advanced/instantiate_objects/overview" + ), + ), + "_recursive_": PropertySchema( + type="boolean", + title="Recursive", + description=( + "Whether instantiating this config should recursively instantiate children configs.\n" + "See: https://hydra.cc/docs/advanced/instantiate_objects/overview/#recursive-instantiation" + ), + ), + }, + dependentRequired={ + "_convert_": ["_target_"], + "_partial_": ["_target_"], + "_args_": ["_target_"], + "_recursive_": ["_target_"], + }, +) + + +def main(argv: list[str] | None = None): + logging.basicConfig( + level=logging.ERROR, + # format="%(asctime)s - %(levelname)s - %(message)s", + format="%(message)s", + datefmt="[%X]", + force=True, + handlers=[ + rich.logging.RichHandler( + markup=True, + rich_tracebacks=True, + tracebacks_width=100, + tracebacks_show_locals=False, + ) + ], + ) + + from project.main import PROJECT_NAME + from project.utils.env_vars import REPO_ROOTDIR + + CONFIGS_DIR = REPO_ROOTDIR / PROJECT_NAME / "configs" + + # FIXME: remove this, perhaps it could be an argument? + from project.utils.env_vars import REPO_ROOTDIR + + repo_root = REPO_ROOTDIR + + parser = argparse.ArgumentParser() + + parser.add_argument( + "path", + type=Path, + default=REPO_ROOTDIR, + help="Directory containing configs, of path to config file.", + ) + parser.add_argument("--configs-dir", type=Path, default=CONFIGS_DIR) + parser.add_argument("--schemas-dir", type=Path, default=None, required=False) + parser.add_argument("--regen-schemas", action=argparse.BooleanOptionalAction) + parser.add_argument("--stop-on-error", action=argparse.BooleanOptionalAction) + parser.add_argument( + "--add-headers", + action=argparse.BooleanOptionalAction, + help=( + "Always add headers to the yaml config files, instead of the default " + "behaviour which is to first try to add an entry in the vscode " + "settings.json file." + ), + ) + verbosity_group = parser.add_mutually_exclusive_group() + verbosity_group.add_argument( + "-q", "--quiet", dest="quiet", action=argparse.BooleanOptionalAction + ) + verbosity_group.add_argument("-v", "--verbose", dest="verbose", action="count", default=0) + + args = parser.parse_args(argv if argv is not None else sys.argv[1:]) + + path: Path = args.path + configs_dir: Path = args.configs_dir + schemas_dir: Path = args.schemas_dir + regen_schemas: bool = args.regen_schemas + stop_on_error: bool = args.stop_on_error + quiet: bool = args.quiet + verbose: int = args.verbose + add_headers: bool = args.add_headers + + if quiet: + logger.setLevel(logging.NOTSET) + elif verbose: + if verbose >= 3: + logger.setLevel(logging.DEBUG) + elif verbose == 2: + logger.setLevel(logging.INFO) + else: + assert verbose == 1 + logger.setLevel(logging.WARNING) + else: + logger.setLevel(logging.ERROR) + + if not path.is_absolute(): + path = Path.cwd() / path + if path.is_dir(): + config_files = _yaml_files_in(configs_dir) + else: + config_files = [path] + + add_schemas_to_all_hydra_configs( + config_files=config_files, + repo_root=repo_root, + configs_dir=configs_dir, + schemas_dir=schemas_dir, + regen_schemas=regen_schemas, + stop_on_error=stop_on_error, + quiet=quiet, + add_headers=add_headers, + ) + logger.info("Done updating the schemas for the Hydra config files.") + + +def _yaml_files_in(configs_dir): + return list(configs_dir.rglob("*.yaml")) + list(configs_dir.rglob("*.yml")) + + +def add_schemas_to_all_hydra_configs( + config_files: list[Path] | None, + repo_root: Path, + configs_dir: Path, + schemas_dir: Path | None = None, + regen_schemas: bool = False, + stop_on_error: bool = False, + quiet: bool = False, + add_headers: bool | None = None, +): + if config_files is None: + config_files = _yaml_files_in(configs_dir) + if not config_files: + warnings.warn(RuntimeWarning("No config files were passed. Skipping.")) + return + + if schemas_dir is None: + schemas_dir = repo_root / ".schemas" + + if schemas_dir.is_relative_to(repo_root): + _add_schemas_dir_to_gitignore(schemas_dir, repo_root=repo_root) + + config_file_to_schema_file: dict[Path, Path] = {} + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=tqdm.TqdmExperimentalWarning) + pbar = tqdm_rich( + config_files, + desc="Creating schemas for Hydra config files...", + total=len(config_files), + leave=not quiet, + ) + + for config_file in pbar: + pretty_config_file_name = config_file.relative_to(configs_dir) + schema_file = _get_schema_file_path(config_file, schemas_dir) + + # todo: check the modification time. If the config file was modified after the schema file, regen the schema file. + + if schema_file.exists(): + schema_file_modified_time = datetime.datetime.fromtimestamp( + schema_file.stat().st_mtime + ) + config_file_modified_time = datetime.datetime.fromtimestamp( + config_file.stat().st_mtime + ) + # Add a delay to account for the time it takes to modify the the config file at the end to remove a header. + if (config_file_modified_time - schema_file_modified_time) > datetime.timedelta( + seconds=10 + ): + logger.info( + f"Config file {pretty_config_file_name} was modified, regenerating the schema." + ) + elif regen_schemas: + pass # regenerate it. + elif _is_incomplete_schema(schema_file): + logger.info( + f"Unable to properly create the schema for {pretty_config_file_name} last time. Trying again." + ) + else: + logger.debug( + f"Schema file {_relative_to_cwd(schema_file)} was already successfully created. Skipping." + ) + continue + + pbar.set_postfix_str(f"Creating schema for {pretty_config_file_name}") + + try: + logger.debug(f"Creating a schema for {pretty_config_file_name}") + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=UserWarning) + config = _load_config(config_file, configs_dir=configs_dir) + schema = create_schema_for_config( + config, config_file=config_file, configs_dir=configs_dir + ) + schema_file.parent.mkdir(exist_ok=True, parents=True) + schema_file.write_text(json.dumps(schema, indent=2).rstrip() + "\n\n") + _set_is_incomplete_schema(schema_file, False) + except ( + pydantic.errors.PydanticSchemaGenerationError, + hydra.errors.MissingConfigException, + hydra.errors.ConfigCompositionException, + omegaconf.errors.InterpolationResolutionError, + ) as exc: + logger.warning( + f"Unable to create a schema for config {pretty_config_file_name}: {exc}" + ) + if stop_on_error: + raise + + schema = copy.deepcopy(HYDRA_CONFIG_SCHEMA) + schema["additionalProperties"] = True + schema["title"] = f"Partial schema for {pretty_config_file_name}" + schema["description"] = ( + f"(errors occurred while trying to create the schema from the signature:\n{exc}" + ) + schema_file.write_text(json.dumps(schema, indent=2) + "\n") + _set_is_incomplete_schema(schema_file, True) + + config_file_to_schema_file[config_file] = schema_file + + # Option 1: Add a vscode setting that associates the schema file with the yaml files. (less intrusive perhaps). + # Option 2: Add a header to the yaml files that points to the schema file. + + # We will use option 1 if a `code` executable is found. + set_schemas_in_vscode_settings_file = bool(shutil.which("code")) + + if add_headers: + set_schemas_in_vscode_settings_file = False + + if set_schemas_in_vscode_settings_file: + try: + logger.debug( + "Found the `code` executable, will add schema paths to the vscode settings." + ) + _install_yaml_vscode_extension() + _add_schemas_to_vscode_settings(config_file_to_schema_file, repo_root=repo_root) + return + except Exception as exc: + logger.error( + f"Unable to write schemas in the vscode settings file. " + f"Falling back to adding a header to config files. (exc={exc})" + ) + + logger.debug("A headers to config files to point to the schemas to use.") + for config_file, schema_file in config_file_to_schema_file.items(): + add_schema_header(config_file, schema_path=schema_file) + + +def _add_schemas_dir_to_gitignore(schemas_dir: Path, repo_root: Path): + _rel = schemas_dir.relative_to(repo_root) + _gitignore_file = repo_root / ".gitignore" + if not any(line.startswith(str(_rel)) for line in _gitignore_file.read_text().splitlines()): + logger.info(f"Adding entry in .gitignore for the schemas directory ({schemas_dir})") + with _gitignore_file.open("a") as f: + f.write(f"{_rel}\n") + + +_incomplete_schema_xattr = "user.schema_error" + + +def _is_incomplete_schema(schema_file: Path) -> bool: + try: + return os.getxattr(schema_file, _incomplete_schema_xattr) == bytes(1) + except OSError: + return False + + +def _set_is_incomplete_schema(schema_file: Path, val: bool): + os.setxattr(schema_file, _incomplete_schema_xattr, bytes(val)) + + +def _relative_to_cwd(p: str | Path): + return Path(p).relative_to(Path.cwd()) + + +def _install_yaml_vscode_extension(): + logger.debug( + "Running `code --install-extension redhat.vscode-yaml` to install the yaml extension for vscode." + ) + output = subprocess.check_output( + ("code", "--install-extension", "redhat.vscode-yaml"), text=True + ) + logger.debug(output) + + +def _add_schemas_to_vscode_settings( + config_file_to_schema_file: dict[Path, Path], + repo_root: Path, +) -> None: + # Make the vscode settings file if necessary: + vscode_dir = repo_root / ".vscode" + vscode_dir.mkdir(exist_ok=True, parents=False) + vscode_settings_file = vscode_dir / "settings.json" + vscode_settings_file.touch(exist_ok=True) + + # TODO: Use Jsonc to load the file and preserve comments? + logger.debug(f"Reading the VsCode settings file at {vscode_settings_file}.") + vscode_settings_content = vscode_settings_file.read_text() + # Remove any trailing commas from the content: + vscode_settings_content = ( + vscode_settings_content.strip().removesuffix("}").rstrip().rstrip(",") + "}" + ) + vscode_settings: dict[str, Any] = json.loads(vscode_settings_content) + + # Avoid the popup and do users a favour by disabling telemetry. + vscode_settings.setdefault("redhat.telemetry.enabled", False) + + yaml_schemas_setting: dict[str, str | list[str]] = vscode_settings.setdefault( + "yaml.schemas", {} + ) + + # Write all the schemas + for config_file, schema_file in config_file_to_schema_file.items(): + assert schema_file.exists() + + _root = vscode_settings_file.parent.parent + schema_key = str(schema_file.relative_to(_root)) + path_to_add = str(config_file.absolute()) + + if schema_key not in yaml_schemas_setting: + yaml_schemas_setting[schema_key] = path_to_add + elif isinstance(files_associated_with_schema := yaml_schemas_setting[schema_key], str): + files = sorted(set([files_associated_with_schema, path_to_add])) + yaml_schemas_setting[schema_key] = files[0] if len(files) == 1 else files + else: + files = sorted(set(files_associated_with_schema + [path_to_add])) + yaml_schemas_setting[schema_key] = files[0] if len(files) == 1 else files + + vscode_settings_file.write_text(json.dumps(vscode_settings, indent=2)) + logger.info(f"Updated the yaml schemas in the vscode settings file at {vscode_settings_file}.") + + # If this worked, then remove any schema directives from the config files. + for config_file, schema_file in config_file_to_schema_file.items(): + assert schema_file.exists() + config_lines = config_file.read_text().splitlines() + lines_to_remove: list[int] = [] + for i, line in enumerate(config_lines): + if line.strip().startswith("# yaml-language-server: $schema="): + lines_to_remove.append(i) + for line_to_remove in reversed(lines_to_remove): + config_lines.pop(line_to_remove) + config_file.write_text("\n".join(config_lines).rstrip() + ("\n" if config_lines else "")) + + +def _get_schema_file_path(config_file: Path, schemas_dir: Path): + config_group = config_file.parent + schema_file = schemas_dir / f"{config_group.name}_{config_file.stem}_schema.json" + return schema_file + + +def create_schema_for_config( + config: dict | DictConfig, config_file: Path, configs_dir: Path | None +) -> Schema | ObjectSchema: + """IDEA: Create a schema for the given config. + + - If you encounter a key, add it to the schema. + - If you encounter a value with a _target_, use a dedicated function to get the schema for that target, and merge it into the current schema. + - Only the top-level config (`config`) can have a `defaults: list[str]` key. + - Should ideally load the defaults and merge this schema on top of them. + """ + schema = copy.deepcopy(HYDRA_CONFIG_SCHEMA) + pretty_path = config_file.relative_to(configs_dir) if configs_dir else config_file + schema["title"] = f"Auto-generated schema for {pretty_path}" + + if config_file.exists() and configs_dir: + # note: the `defaults` list gets consumed by Hydra in `_load_config`, so we actually re-read the + # config file to get the `defaults`, if present. + # TODO: There still defaults here, only during unit tests, even though _load_config should have consumed them! + if "defaults" in config: + schema = _update_schema_from_defaults( + config_file=config_file, + schema=schema, + defaults=config["defaults"], + configs_dir=configs_dir, + ) + + if target_name := config.get("_target_"): + # There's a '_target_' key at the top level in the config file. + target = hydra.utils.get_object(target_name) + schema["description"] = f"Based on the signature of {target}." + if "properties" not in schema: + schema["properties"] = {} + assert "properties" in schema and isinstance(schema["properties"], dict) + schema["properties"]["_target_"] = PropertySchema( + type="string", + title="Target", + const=target_name, + # pattern=r"", # todo: Use a pattern to match python module import strings. + description=( + f"Target to instantiate, in this case: `{target_name}`\n" + # f"* Source: \n" + # f"* Config file: \n" + f"See the Hydra docs for '_target_': https://hydra.cc/docs/advanced/instantiate_objects/overview/\n" + ), + ) + + nested_value_schema_from_target_signature = _get_schema_from_target(config) + # logger.debug(f"Schema from signature of {target}: {schema_from_target_signature}") + + schema = _merge_dicts( + nested_value_schema_from_target_signature, # type: ignore + schema, # type: ignore + conflict_handler=overwrite, + ) + + return schema + + # Config file that contains entries that may or may not have a _target_. + schema["additionalProperties"] = True + + def all_subentries_with_target(config: dict) -> dict[tuple[str, ...], dict]: + """Iterator that yields all the nested config entries that have a _target_.""" + entries = {} + for key, value in config.items(): + if isinstance(value, dict) and "_target_" in value.keys(): + entries[(key,)] = value + elif isinstance(value, dict): + for subkey, subvalue in all_subentries_with_target(value).items(): + entries[(key, *subkey)] = subvalue + return entries + + _config_dict = ( + OmegaConf.to_container(config, resolve=False) if isinstance(config, DictConfig) else config + ) + assert isinstance(_config_dict, dict) + for keys, value in all_subentries_with_target(_config_dict).items(): + # Go over all the values in the config. If any of them have a `_target_`, then we can + # add a schema at that entry. + assert "_target_" in value + target = hydra.utils.get_object(value["_target_"]) + + # try: + nested_value_schema_from_target_signature = _get_schema_from_target(value) + # except omegaconf.errors.InterpolationToMissingValueError: + # logger.warning( + # f"Unable to get the schema for {value['_target_']} at key {'.'.join(keys)} " + # f"in file {config_file}." + # ) + # continue + + if "$defs" in nested_value_schema_from_target_signature: + # note: can't have a $defs key in the schema. + schema.setdefault("$defs", {}).update( # type: ignore + nested_value_schema_from_target_signature.pop("$defs") + ) + + logger.debug( + f"Getting schema from target {value['_target_']} at key {'.'.join(keys)} in " + f"{config_file}." + ) + + assert "properties" in nested_value_schema_from_target_signature + + parent_keys, last_key = keys[:-1], keys[-1] + where_to_set: Schema | ObjectSchema = schema + for key in parent_keys: + where_to_set = where_to_set.setdefault("properties", {}).setdefault(key, {}) # type: ignore + + if "properties" not in where_to_set: + where_to_set["properties"] = {last_key: nested_value_schema_from_target_signature} # type: ignore + elif last_key not in where_to_set["properties"]: + assert isinstance(last_key, str) + where_to_set["properties"][last_key] = nested_value_schema_from_target_signature # type: ignore + else: + where_to_set["properties"] = _merge_dicts( # type: ignore + where_to_set["properties"], + nested_value_schema_from_target_signature, # type: ignore + ) + raise NotImplementedError("todo: use merge_dicts here") + + return schema + + +def _update_schema_from_defaults( + config_file: Path, + schema: Schema, + defaults: list[str | dict[str, str]], + configs_dir: Path, +): + defaults_list = defaults + + for default in defaults_list: + if default == "_self_": # todo: does this actually make sense? + continue + # Note: The defaults can also have the .yaml or .yml extension, _load_config drops the + # extension. + if isinstance(default, str): + assert not default.startswith("/") + other_config_path = config_file.parent / default + else: + assert len(default) == 1 + key, val = next(iter(default.items())) + other_config_path = config_file.parent / key / val + logger.debug(f"Loading config of default {default}.") + + if not other_config_path.suffix: + # If the other config file has the name without the extension, try both .yml and .yaml. + for suffix in (".yml", ".yaml"): + if other_config_path.with_suffix(suffix).exists(): + other_config_path = other_config_path.with_suffix(suffix) + break + + # try: + default_config = _load_config(other_config_path, configs_dir=configs_dir) + # except omegaconf.errors.MissingMandatoryValue: + # default_config = OmegaConf.load(other_config_path) + + schema_of_default = create_schema_for_config( + config=default_config, + config_file=other_config_path, + configs_dir=configs_dir, + ) + + logger.debug(f"Schema from default {default}: {schema_of_default}") + logger.debug(f"Properties of {default=}: {list(schema_of_default['properties'].keys())}") # type: ignore + + schema = _merge_dicts( # type: ignore + schema_of_default, # type: ignore + schema, # type: ignore + conflict_handlers={ + "_target_": overwrite, # use the new target. + "default": overwrite, # use the new default? + "title": overwrite, + "description": overwrite, + }, + ) + # todo: deal with this one here. + if schema.get("additionalProperties") is False: + schema.pop("additionalProperties") + return schema + + +def overwrite(val_a: Any, val_b: Any) -> Any: + return val_b + + +def keep_previous(val_a: Any, val_b: Any) -> Any: + return val_a + + +conflict_handlers: dict[str, Callable[[Any, Any], Any]] = {} + +_K = TypeVar("_K") +_V = TypeVar("_V") +_NestedDict = MutableMapping[_K, _V | "_NestedDict[_K, _V]"] + +_D1 = TypeVar("_D1", bound=_NestedDict) +_D2 = TypeVar("_D2", bound=_NestedDict) + + +def _merge_dicts( + a: _D1, + b: _D2, + path: list[str] = [], + conflict_handlers: dict[str, Callable[[Any, Any], Any]] = conflict_handlers, + conflict_handler: Callable[[Any, Any], Any] | None = None, +) -> _D1 | _D2: + """Merge two nested dictionaries. + + >>> x = dict(b=1, c=dict(d=2, e=3)) + >>> y = dict(d=3, c=dict(z=2, f=4)) + >>> _merge_dicts(x, y) + {'b': 1, 'c': {'d': 2, 'e': 3, 'z': 2, 'f': 4}, 'd': 3} + >>> x + {'b': 1, 'c': {'d': 2, 'e': 3}} + >>> y + {'d': 3, 'c': {'z': 2, 'f': 4}} + """ + out = copy.deepcopy(a) + for key in b: + if key in a: + if isinstance(a[key], dict) and isinstance(b[key], dict): + out[key] = _merge_dicts( + a[key], + b[key], + path + [str(key)], + conflict_handlers={ + k.removeprefix(f"{key}."): v for k, v in conflict_handlers.items() + }, + conflict_handler=conflict_handler, + ) + elif a[key] != b[key]: + if specific_conflict_handler := conflict_handlers.get(key): + out[key] = specific_conflict_handler(a[key], b[key]) # type: ignore + elif conflict_handler: + out[key] = conflict_handler(a[key], b[key]) # type: ignore + + # if any(key.split(".")[-1] == handler_name for for prefix in ["_", "description", "title"]): + # out[key] = b[key] + else: + raise Exception("Conflict at " + ".".join(path + [str(key)])) + else: + out[key] = copy.deepcopy(b[key]) # type: ignore + return out + + +def _has_package_global_line(config_file: Path) -> int | None: + """Returns whether the config file contains a `@package _global_` directive of hydra. + + See: https://hydra.cc/docs/advanced/overriding_packages/#overriding-the-package-via-the-package-directive + """ + for line in config_file.read_text().splitlines(): + line = line.strip() + if not line.startswith("#"): + continue + if line.removeprefix("#").strip().startswith("@package _global_"): + return True + return False + + +def _load_config(config_path: Path, configs_dir: Path) -> DictConfig: + *config_groups, config_name = config_path.relative_to(configs_dir).with_suffix("").parts + logger.debug( + f"config_path: ./{_relative_to_cwd(config_path)}, {config_groups=}, {config_name=}, configs_dir: {configs_dir}" + ) + config_group = "/".join(config_groups) + + from hydra.core.utils import setup_globals + + # todo: When would this normally be called? + setup_globals() + + # FIXME! + if configs_dir.is_relative_to(REPO_ROOTDIR) and (configs_dir / "__init__.py").exists(): + config_module = str(configs_dir.relative_to(REPO_ROOTDIR)).replace("/", ".") + search_path = create_config_search_path(f"pkg://{config_module}") + else: + search_path = create_config_search_path(str(configs_dir)) + + if _has_package_global_line(config_path): + # Tricky: Here we load the global config but with the given config as an override. + config_loader = ConfigLoaderImpl(config_search_path=search_path) + top_config = config_loader.load_configuration( + "config", # todo: Fix this? + overrides=[f"{config_group}={config_name}"], + # todo: setting this here because it appears to be what's used in Hydra in a normal + # run, even though RunMode.RUN would make more sense intuitively. + run_mode=RunMode.MULTIRUN, + ) + return top_config + + # Load the global config and get the node for the desired config. + # TODO: Can this cause errors if configs in an unrelated subtree have required values? + logger.debug(f"loading config {config_path}") + config_loader = ConfigLoaderImpl(config_search_path=search_path) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=UserWarning) + top_config = config_loader.load_configuration( + f"{config_group}/{config_name}", overrides=[], run_mode=RunMode.MULTIRUN + ) + # Retrieve the sub-entry in the config and return it. + config = top_config + for config_group in config_groups: + config = config[config_group] + return config + + +def add_schema_header(config_file: Path, schema_path: Path) -> None: + """Add a comment in the yaml config file to tell yaml language server where to look for the + schema. + + Importantly in the context of Hydra, this comment line should be added **after** any `# + @package: ` directives of Hydra, otherwise Hydra doesn't use those directives properly + anymore. + """ + lines = config_file.read_text().splitlines(keepends=False) + + if config_file.parent is schema_path.parent: + # TODO: Unsure when this branch would be used, and if it would differ. + relative_path_to_schema = "./" + schema_path.name + else: + relative_path_to_schema = os.path.relpath(schema_path, start=config_file.parent) + + # Remove any existing schema lines. + lines = [ + line for line in lines if not line.strip().startswith("# yaml-language-server: $schema=") + ] + + # NOTE: This line can be placed anywhere in the file, not necessarily needs to be at the top, + # and the yaml vscode extension will pick it up. + new_line = f"# yaml-language-server: $schema={relative_path_to_schema}" + + package_global_line: int | None = None + + for i, line in enumerate(lines): + line = line.strip() + if not line: + continue + # BUG: IF the schema line comes before a @package: global comment, then the @package: _global_ + # comment is ignored by Hydra. + # Locate the last package line (a bit unnecessary, since there should only be one). + if line.startswith("#") and line.removeprefix("#").strip().startswith("@package"): + package_global_line = i + + if package_global_line is None: + # There's no package directive in the file. + new_lines = [new_line, *lines] + else: + new_lines = lines.copy() + # Insert the schema line after the package directive line. + new_lines.insert(package_global_line + 1, new_line) + + result = "\n".join(new_lines).strip() + "\n" + if config_file.read_text() != result: + config_file.write_text(result) + + +def _get_schema_from_target(config: dict | DictConfig) -> ObjectSchema | Schema: + assert isinstance(config, dict | DictConfig) + logger.debug(f"Config: {config}") + target = hydra.utils.get_object(config["_target_"]) + + if inspect.isclass(target) and issubclass(target, flax.linen.Module): + object_type = hydra_zen.builds( + target, + populate_full_signature=True, + hydra_recursive=False, + hydra_convert="all", + zen_exclude=["parent"], + zen_dataclass={"cls_name": target.__qualname__}, + ) + + elif inspect.isclass(target) and issubclass( + target, lightning.pytorch.callbacks.RichProgressBar + ): + from rich.style import Style # noqa + + # todo: trying to fix this here. + # RichProgressBarTheme.__annotations__[] + object_type = hydra_zen.builds( + target, + populate_full_signature=True, + hydra_recursive=False, + hydra_convert="all", + zen_exclude=["theme"], + zen_dataclass={"cls_name": target.__qualname__}, + ) + elif dataclasses.is_dataclass(target): + # The target is a dataclass, so the schema is just the schema of the dataclass. + object_type = target + else: + # The target is a type or callable. + assert callable(target) + object_type = hydra_zen.builds( + target, + populate_full_signature=True, + hydra_defaults=config.get("defaults", None), + hydra_recursive=False, + hydra_convert="all", + zen_dataclass={"cls_name": target.__qualname__}, + # zen_wrappers=pydantic_parser, # unsure if this is how it works? + ) + + try: + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=UserWarning) + + json_schema = pydantic.TypeAdapter(object_type).json_schema( + mode="serialization", + schema_generator=_MyGenerateJsonSchema, + by_alias=False, + ) + json_schema = typing.cast(Schema, json_schema) + + # if "$defs" in json_schema: + # for key, class_schema in list(json_schema["$defs"].items()): + # logger.debug(f"Before resolving {key}: {json_schema}") + # json_schema["$defs"].pop(key) + # class_schema["title"] = key + # expanded = json.dumps(json_schema).replace( + # json.dumps({"$ref": f"#/$defs/{key}"}), json.dumps(class_schema) + # ) + # logger.debug(f"Expanded: {expanded}") + # json_schema = json.loads(expanded) + # logger.debug(f"After resolving {key}: {json_schema}") + # assert False, json_schema + assert "properties" in json_schema + except pydantic.PydanticSchemaGenerationError as e: + raise NotImplementedError(f"Unable to get the schema with pydantic: {e}") + + assert "properties" in json_schema + + docs_to_search: list[dp.Docstring] = [] + + if inspect.isclass(target): + for target_or_base_class in inspect.getmro(target): + if class_docstring := inspect.getdoc(target_or_base_class): + docs_to_search.append(dp.parse(class_docstring)) + if init_docstring := inspect.getdoc(target_or_base_class.__init__): + docs_to_search.append(dp.parse(init_docstring)) + else: + assert inspect.isfunction(target) + docstring = inspect.getdoc(target) + if docstring: + docs_to_search = [dp.parse(docstring)] + + param_descriptions: dict[str, str] = {} + for doc in docs_to_search: + for param in doc.params: + if param.description and param.arg_name not in param_descriptions: + param_descriptions[param.arg_name] = param.description + + # Update the pydantic schema with descriptions: + for property_name, property_dict in json_schema["properties"].items(): + if description := param_descriptions.get(property_name): + property_dict["description"] = description + else: + property_dict["description"] = ( + f"The {property_name} parameter of the {target.__qualname__}." + ) + + if config.get("_partial_"): + json_schema["required"] = [] + + # if the target takes **kwargs, then we don't restrict additional properties. + json_schema["additionalProperties"] = inspect.getfullargspec(target).varkw is not None + + return json_schema + + +def _target_has_var_kwargs(config: DictConfig) -> bool: + target = hydra_zen.get_target(config) # type: ignore + return inspect.getfullargspec(target).varkw is None + + +class _MyGenerateJsonSchema(GenerateJsonSchema): + # def handle_invalid_for_json_schema( + # self, schema: core_schema.CoreSchema, error_info: str + # ) -> JsonSchemaValue: + # raise PydanticOmit + + def enum_schema(self, schema: core_schema.EnumSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches an Enum value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + enum_type = schema["cls"] + logger.debug(f"Enum of type {enum_type}") + import torchvision.models.resnet + + if issubclass(enum_type, torchvision.models.WeightsEnum): + + @dataclasses.dataclass + class Dummy: + value: str + + slightly_changed_schema = schema | { + "members": [Dummy(v.name) for v in schema["members"]] + } + return super().enum_schema(slightly_changed_schema) # type: ignore + return super().enum_schema(schema) + + +class AutoSchemaPlugin(SearchPathPlugin): + provider: str + path: str + + def __init__(self) -> None: + super().__init__() + + def manipulate_search_path(self, search_path: ConfigSearchPath) -> None: + search_path_entries = search_path.get_path() + # WIP: Trying to infer the project root, configs dir from the Hydra context. + # Currently hard-coded. + + for seach_path_entry in search_path_entries: + if seach_path_entry.provider != "main": + continue + + if seach_path_entry.path.startswith("pkg://"): + configs_pkg = seach_path_entry.path.removeprefix("pkg://") + project_package = configs_pkg.split(".")[0] + _project_module = importlib.import_module(project_package) + assert ( + _project_module.__file__ + and Path(_project_module.__file__).name == "__init__.py" + ) + repo_root = Path(_project_module.__file__).parent.parent + configs_dir = repo_root / configs_pkg.replace(".", "/") + if not (repo_root.is_dir() and configs_dir.is_dir()): + logger.warning( + f"Unable to add schemas to Hydra configs: " + f"Expected to find the project root directory at {repo_root} " + f"and the configs directory at {configs_dir}!" + ) + continue + logger.debug("Assuming that the ") + add_schemas_to_all_hydra_configs( + config_files=None, + repo_root=repo_root, + configs_dir=configs_dir, + quiet=True, + ) + + +if __name__ == "__main__": + main() diff --git a/project/utils/auto_schema_test.py b/project/utils/auto_schema_test.py new file mode 100644 index 00000000..6484f478 --- /dev/null +++ b/project/utils/auto_schema_test.py @@ -0,0 +1,94 @@ +import json +import subprocess +from pathlib import Path + +import hydra.errors +import pytest +import yaml +from pytest_regressions.file_regression import FileRegressionFixture + +from .auto_schema import add_schema_header, create_schema_for_config, main + + +class Foo: + def __init__(self, bar: str): + """Description of the `Foo` class. + + Args: + bar: Description of the `bar` argument. + """ + self.bar = bar + + +class Bar(Foo): + """Docstring of the Bar class. + + Args: + baz: description of the `baz` argument from the cls docstring instead of the init docstring. + """ + + def __init__(self, bar: str, baz: int): + # no docstring here. + super().__init__(bar=bar) + self.baz = baz + + +_this_file = Path(__file__) +test_files = list((_this_file.parent / _this_file.name).with_suffix("").rglob("*.yaml")) + + +@pytest.mark.parametrize("config_file", test_files, ids=[f.name for f in test_files]) +def test_make_schema( + config_file: Path, file_regression: FileRegressionFixture, original_datadir: Path +): + """Test that creates a schema for a config file and saves it next to it. + + (in the test folder). + """ + schema_file = config_file.with_suffix(".json") + + config = yaml.load(config_file.read_text(), yaml.FullLoader) + schema = create_schema_for_config( + config=config, config_file=config_file, configs_dir=original_datadir + ) + + add_schema_header(config_file, schema_path=schema_file) + file_regression.check( + json.dumps(schema, indent=2) + "\n", fullpath=schema_file, extension=".json" + ) + + +def test_can_run_via_cli(): + """Actually run the command on the repo from the CLI.""" + # Run programmatically instead of with a subproc4ess so we can get nice coverage stats. + main(["."]) # assuming we're at the project root directory. + + +@pytest.mark.xfail( + raises=hydra.errors.MissingConfigException, + reason="TODO: Missing hydra/sweeper/orion config since it isn't installed atm.", + strict=True, +) +def test_run_via_cli_without_errors(): + """Checks that the command completes without errors.""" + # Run programmatically instead of with a subproc4ess so we can get nice coverage stats. + main([".", "--stop-on-error"]) # assuming we're at the project root directory. + + +def test_run_via_rye_script(): + """Actually run the command on the repo, via the `[tool.rye.scripts]` entry in + pyproject.toml.""" + # Run once so we can get nice coverage stats. + subprocess.check_call(["rye", "run", "auto_schema"], text=True) + + +@pytest.mark.xfail( + raises=subprocess.CalledProcessError, + reason="TODO: Missing hydra/sweeper/orion config since it isn't installed atm.", + strict=True, +) +def test_run_via_rye_script_without_errors(): + """Actually run the command on the repo, via the `[tool.rye.scripts]` entry in + pyproject.toml.""" + # Run once so we can get nice coverage stats. + subprocess.check_call(["rye", "run", "auto_schema", "--stop-on-error"], text=True) diff --git a/project/utils/auto_schema_test/nested.json b/project/utils/auto_schema_test/nested.json new file mode 100644 index 00000000..5c92334f --- /dev/null +++ b/project/utils/auto_schema_test/nested.json @@ -0,0 +1,129 @@ +{ + "title": "Auto-generated schema for nested.yaml", + "description": "Schema created by the `auto_schema.py` script.", + "properties": { + "defaults": { + "title": "Hydra defaults", + "description": "Hydra defaults for this config. See https://hydra.cc/docs/advanced/defaults_list/", + "type": "array", + "items": { + "oneOf": [ + { + "type": "object", + "propertyNames": { + "pattern": "^(override\\s*)?(/?\\w*)+$" + }, + "patternProperties": { + "^(override\\s*)?(/?\\w*)*$": { + "type": "string", + "pattern": "\\w*(.yaml|.yml)?$" + } + }, + "minProperties": 1, + "maxProperties": 1 + }, + { + "type": "string", + "pattern": "^\\w+(.yaml|.yml)?$" + }, + { + "type": "object", + "propertyNames": { + "pattern": "^(override\\s*)?(/?\\w*)+$" + }, + "patternProperties": { + "^(override\\s*)?(/?\\w*)*$": { + "type": "null" + } + }, + "minProperties": 1, + "maxProperties": 1 + } + ] + }, + "uniqueItems": true + }, + "_target_": { + "type": "string", + "title": "Target", + "description": "Target to instantiate.\nSee https://hydra.cc/docs/advanced/instantiate_objects/overview/" + }, + "_convert_": { + "type": "string", + "enum": [ + "none", + "partial", + "object", + "all" + ], + "title": "Convert", + "description": "See https://hydra.cc/docs/advanced/instantiate_objects/overview/#parameter-conversion-strategies" + }, + "_partial_": { + "type": "boolean", + "title": "Partial", + "description": "Whether this config calls the target function when instantiated, or creates a `functools.partial` that will call the target.\nSee: https://hydra.cc/docs/advanced/instantiate_objects/overview" + }, + "_recursive_": { + "type": "boolean", + "title": "Recursive", + "description": "Whether instantiating this config should recursively instantiate children configs.\nSee: https://hydra.cc/docs/advanced/instantiate_objects/overview/#recursive-instantiation" + }, + "a": { + "properties": { + "b": { + "properties": { + "foo": { + "properties": { + "_target_": { + "default": "project.utils.auto_schema_test.Foo", + "title": " Target ", + "type": "string", + "description": "The _target_ parameter of the Foo." + }, + "_recursive_": { + "default": false, + "title": " Recursive ", + "type": "boolean", + "description": "The _recursive_ parameter of the Foo." + }, + "_convert_": { + "default": "all", + "title": " Convert ", + "type": "string", + "description": "The _convert_ parameter of the Foo." + }, + "bar": { + "title": "Bar", + "type": "string", + "description": "Description of the `bar` argument." + } + }, + "required": [ + "bar" + ], + "title": "Foo", + "type": "object", + "additionalProperties": false + } + } + } + } + } + }, + "dependentRequired": { + "_convert_": [ + "_target_" + ], + "_partial_": [ + "_target_" + ], + "_args_": [ + "_target_" + ], + "_recursive_": [ + "_target_" + ] + }, + "additionalProperties": true +} diff --git a/project/utils/auto_schema_test/nested.yaml b/project/utils/auto_schema_test/nested.yaml new file mode 100644 index 00000000..13baa9a0 --- /dev/null +++ b/project/utils/auto_schema_test/nested.yaml @@ -0,0 +1,6 @@ +# yaml-language-server: $schema=nested.json +a: + b: + foo: + _target_: project.utils.auto_schema_test.Foo + bar: "bob" diff --git a/project/utils/auto_schema_test/partial.json b/project/utils/auto_schema_test/partial.json new file mode 100644 index 00000000..2bfe5fa0 --- /dev/null +++ b/project/utils/auto_schema_test/partial.json @@ -0,0 +1,99 @@ +{ + "properties": { + "_target_": { + "default": "project.utils.auto_schema_test.Foo", + "title": "Target", + "type": "string", + "description": "Target to instantiate, in this case: `project.utils.auto_schema_test.Foo`\nSee the Hydra docs for '_target_': https://hydra.cc/docs/advanced/instantiate_objects/overview/\n", + "const": "project.utils.auto_schema_test.Foo" + }, + "_recursive_": { + "default": false, + "title": "Recursive", + "type": "boolean", + "description": "Whether instantiating this config should recursively instantiate children configs.\nSee: https://hydra.cc/docs/advanced/instantiate_objects/overview/#recursive-instantiation" + }, + "_convert_": { + "default": "all", + "title": "Convert", + "type": "string", + "description": "See https://hydra.cc/docs/advanced/instantiate_objects/overview/#parameter-conversion-strategies", + "enum": [ + "none", + "partial", + "object", + "all" + ] + }, + "bar": { + "title": "Bar", + "type": "string", + "description": "Description of the `bar` argument." + }, + "defaults": { + "title": "Hydra defaults", + "description": "Hydra defaults for this config. See https://hydra.cc/docs/advanced/defaults_list/", + "type": "array", + "items": { + "oneOf": [ + { + "type": "object", + "propertyNames": { + "pattern": "^(override\\s*)?(/?\\w*)+$" + }, + "patternProperties": { + "^(override\\s*)?(/?\\w*)*$": { + "type": "string", + "pattern": "\\w*(.yaml|.yml)?$" + } + }, + "minProperties": 1, + "maxProperties": 1 + }, + { + "type": "string", + "pattern": "^\\w+(.yaml|.yml)?$" + }, + { + "type": "object", + "propertyNames": { + "pattern": "^(override\\s*)?(/?\\w*)+$" + }, + "patternProperties": { + "^(override\\s*)?(/?\\w*)*$": { + "type": "null" + } + }, + "minProperties": 1, + "maxProperties": 1 + } + ] + }, + "uniqueItems": true + }, + "_partial_": { + "type": "boolean", + "title": "Partial", + "description": "Whether this config calls the target function when instantiated, or creates a `functools.partial` that will call the target.\nSee: https://hydra.cc/docs/advanced/instantiate_objects/overview" + } + }, + "required": [], + "title": "Auto-generated schema for partial.yaml", + "type": "object", + "additionalProperties": false, + "description": "Based on the signature of .", + "dependentRequired": { + "_convert_": [ + "_target_" + ], + "_partial_": [ + "_target_" + ], + "_args_": [ + "_target_" + ], + "_recursive_": [ + "_target_" + ] + } +} diff --git a/project/utils/auto_schema_test/partial.yaml b/project/utils/auto_schema_test/partial.yaml new file mode 100644 index 00000000..8970d618 --- /dev/null +++ b/project/utils/auto_schema_test/partial.yaml @@ -0,0 +1,3 @@ +# yaml-language-server: $schema=partial.json +_target_: project.utils.auto_schema_test.Foo +_partial_: true diff --git a/project/utils/auto_schema_test/with_defaults.json b/project/utils/auto_schema_test/with_defaults.json new file mode 100644 index 00000000..cb83119c --- /dev/null +++ b/project/utils/auto_schema_test/with_defaults.json @@ -0,0 +1,101 @@ +{ + "properties": { + "_target_": { + "default": "project.utils.auto_schema_test.Foo", + "title": "Target", + "type": "string", + "description": "Target to instantiate.\nSee https://hydra.cc/docs/advanced/instantiate_objects/overview/", + "const": "project.utils.auto_schema_test.Foo" + }, + "_recursive_": { + "default": false, + "title": "Recursive", + "type": "boolean", + "description": "Whether instantiating this config should recursively instantiate children configs.\nSee: https://hydra.cc/docs/advanced/instantiate_objects/overview/#recursive-instantiation" + }, + "_convert_": { + "default": "all", + "title": "Convert", + "type": "string", + "description": "See https://hydra.cc/docs/advanced/instantiate_objects/overview/#parameter-conversion-strategies", + "enum": [ + "none", + "partial", + "object", + "all" + ] + }, + "bar": { + "title": "Bar", + "type": "string", + "description": "Description of the `bar` argument." + }, + "defaults": { + "title": "Hydra defaults", + "description": "Hydra defaults for this config. See https://hydra.cc/docs/advanced/defaults_list/", + "type": "array", + "items": { + "oneOf": [ + { + "type": "object", + "propertyNames": { + "pattern": "^(override\\s*)?(/?\\w*)+$" + }, + "patternProperties": { + "^(override\\s*)?(/?\\w*)*$": { + "type": "string", + "pattern": "\\w*(.yaml|.yml)?$" + } + }, + "minProperties": 1, + "maxProperties": 1 + }, + { + "type": "string", + "pattern": "^\\w+(.yaml|.yml)?$" + }, + { + "type": "object", + "propertyNames": { + "pattern": "^(override\\s*)?(/?\\w*)+$" + }, + "patternProperties": { + "^(override\\s*)?(/?\\w*)*$": { + "type": "null" + } + }, + "minProperties": 1, + "maxProperties": 1 + } + ] + }, + "uniqueItems": true + }, + "_partial_": { + "type": "boolean", + "title": "Partial", + "description": "Whether this config calls the target function when instantiated, or creates a `functools.partial` that will call the target.\nSee: https://hydra.cc/docs/advanced/instantiate_objects/overview" + } + }, + "required": [ + "bar" + ], + "title": "Auto-generated schema for with_defaults.yaml", + "type": "object", + "description": "Schema created by the `auto_schema.py` script.", + "dependentRequired": { + "_convert_": [ + "_target_" + ], + "_partial_": [ + "_target_" + ], + "_args_": [ + "_target_" + ], + "_recursive_": [ + "_target_" + ] + }, + "additionalProperties": true +} diff --git a/project/utils/auto_schema_test/with_defaults.yaml b/project/utils/auto_schema_test/with_defaults.yaml new file mode 100644 index 00000000..6c6ede0e --- /dev/null +++ b/project/utils/auto_schema_test/with_defaults.yaml @@ -0,0 +1,4 @@ +# yaml-language-server: $schema=with_defaults.json +defaults: + - with_target +bar: "bob" diff --git a/project/utils/auto_schema_test/with_package_global.json b/project/utils/auto_schema_test/with_package_global.json new file mode 100644 index 00000000..d41a2447 --- /dev/null +++ b/project/utils/auto_schema_test/with_package_global.json @@ -0,0 +1,101 @@ +{ + "properties": { + "_target_": { + "default": "project.utils.auto_schema_test.Foo", + "title": "Target", + "type": "string", + "description": "Target to instantiate, in this case: `project.utils.auto_schema_test.Foo`\nSee the Hydra docs for '_target_': https://hydra.cc/docs/advanced/instantiate_objects/overview/\n", + "const": "project.utils.auto_schema_test.Foo" + }, + "_recursive_": { + "default": false, + "title": "Recursive", + "type": "boolean", + "description": "Whether instantiating this config should recursively instantiate children configs.\nSee: https://hydra.cc/docs/advanced/instantiate_objects/overview/#recursive-instantiation" + }, + "_convert_": { + "default": "all", + "title": "Convert", + "type": "string", + "description": "See https://hydra.cc/docs/advanced/instantiate_objects/overview/#parameter-conversion-strategies", + "enum": [ + "none", + "partial", + "object", + "all" + ] + }, + "bar": { + "title": "Bar", + "type": "string", + "description": "Description of the `bar` argument." + }, + "defaults": { + "title": "Hydra defaults", + "description": "Hydra defaults for this config. See https://hydra.cc/docs/advanced/defaults_list/", + "type": "array", + "items": { + "oneOf": [ + { + "type": "object", + "propertyNames": { + "pattern": "^(override\\s*)?(/?\\w*)+$" + }, + "patternProperties": { + "^(override\\s*)?(/?\\w*)*$": { + "type": "string", + "pattern": "\\w*(.yaml|.yml)?$" + } + }, + "minProperties": 1, + "maxProperties": 1 + }, + { + "type": "string", + "pattern": "^\\w+(.yaml|.yml)?$" + }, + { + "type": "object", + "propertyNames": { + "pattern": "^(override\\s*)?(/?\\w*)+$" + }, + "patternProperties": { + "^(override\\s*)?(/?\\w*)*$": { + "type": "null" + } + }, + "minProperties": 1, + "maxProperties": 1 + } + ] + }, + "uniqueItems": true + }, + "_partial_": { + "type": "boolean", + "title": "Partial", + "description": "Whether this config calls the target function when instantiated, or creates a `functools.partial` that will call the target.\nSee: https://hydra.cc/docs/advanced/instantiate_objects/overview" + } + }, + "required": [ + "bar" + ], + "title": "Auto-generated schema for with_package_global.yaml", + "type": "object", + "additionalProperties": false, + "description": "Based on the signature of .", + "dependentRequired": { + "_convert_": [ + "_target_" + ], + "_partial_": [ + "_target_" + ], + "_args_": [ + "_target_" + ], + "_recursive_": [ + "_target_" + ] + } +} diff --git a/project/utils/auto_schema_test/with_package_global.yaml b/project/utils/auto_schema_test/with_package_global.yaml new file mode 100644 index 00000000..d181dc17 --- /dev/null +++ b/project/utils/auto_schema_test/with_package_global.yaml @@ -0,0 +1,4 @@ +# @package: _global_ +# yaml-language-server: $schema=with_package_global.json +_target_: project.utils.auto_schema_test.Foo +bar: "bob" diff --git a/project/utils/auto_schema_test/with_target.json b/project/utils/auto_schema_test/with_target.json new file mode 100644 index 00000000..b75255e6 --- /dev/null +++ b/project/utils/auto_schema_test/with_target.json @@ -0,0 +1,101 @@ +{ + "properties": { + "_target_": { + "default": "project.utils.auto_schema_test.Foo", + "title": "Target", + "type": "string", + "description": "Target to instantiate, in this case: `project.utils.auto_schema_test.Foo`\nSee the Hydra docs for '_target_': https://hydra.cc/docs/advanced/instantiate_objects/overview/\n", + "const": "project.utils.auto_schema_test.Foo" + }, + "_recursive_": { + "default": false, + "title": "Recursive", + "type": "boolean", + "description": "Whether instantiating this config should recursively instantiate children configs.\nSee: https://hydra.cc/docs/advanced/instantiate_objects/overview/#recursive-instantiation" + }, + "_convert_": { + "default": "all", + "title": "Convert", + "type": "string", + "description": "See https://hydra.cc/docs/advanced/instantiate_objects/overview/#parameter-conversion-strategies", + "enum": [ + "none", + "partial", + "object", + "all" + ] + }, + "bar": { + "title": "Bar", + "type": "string", + "description": "Description of the `bar` argument." + }, + "defaults": { + "title": "Hydra defaults", + "description": "Hydra defaults for this config. See https://hydra.cc/docs/advanced/defaults_list/", + "type": "array", + "items": { + "oneOf": [ + { + "type": "object", + "propertyNames": { + "pattern": "^(override\\s*)?(/?\\w*)+$" + }, + "patternProperties": { + "^(override\\s*)?(/?\\w*)*$": { + "type": "string", + "pattern": "\\w*(.yaml|.yml)?$" + } + }, + "minProperties": 1, + "maxProperties": 1 + }, + { + "type": "string", + "pattern": "^\\w+(.yaml|.yml)?$" + }, + { + "type": "object", + "propertyNames": { + "pattern": "^(override\\s*)?(/?\\w*)+$" + }, + "patternProperties": { + "^(override\\s*)?(/?\\w*)*$": { + "type": "null" + } + }, + "minProperties": 1, + "maxProperties": 1 + } + ] + }, + "uniqueItems": true + }, + "_partial_": { + "type": "boolean", + "title": "Partial", + "description": "Whether this config calls the target function when instantiated, or creates a `functools.partial` that will call the target.\nSee: https://hydra.cc/docs/advanced/instantiate_objects/overview" + } + }, + "required": [ + "bar" + ], + "title": "Auto-generated schema for with_target.yaml", + "type": "object", + "additionalProperties": false, + "description": "Based on the signature of .", + "dependentRequired": { + "_convert_": [ + "_target_" + ], + "_partial_": [ + "_target_" + ], + "_args_": [ + "_target_" + ], + "_recursive_": [ + "_target_" + ] + } +} diff --git a/project/utils/auto_schema_test/with_target.yaml b/project/utils/auto_schema_test/with_target.yaml new file mode 100644 index 00000000..35226f42 --- /dev/null +++ b/project/utils/auto_schema_test/with_target.yaml @@ -0,0 +1,3 @@ +# yaml-language-server: $schema=with_target.json +_target_: project.utils.auto_schema_test.Foo +bar: "bob" diff --git a/project/utils/hydra_config_utils.py b/project/utils/hydra_config_utils.py index 3d7fe617..7f1d1094 100644 --- a/project/utils/hydra_config_utils.py +++ b/project/utils/hydra_config_utils.py @@ -16,8 +16,10 @@ def get_config_loader(): from hydra._internal.config_loader_impl import ConfigLoaderImpl from hydra._internal.utils import create_automatic_config_search_path + from project.main import PROJECT_NAME + search_path = create_automatic_config_search_path( - calling_file=None, calling_module=None, config_path="pkg://project.configs" + calling_file=None, calling_module=None, config_path=f"pkg://{PROJECT_NAME}.configs" ) config_loader = ConfigLoaderImpl(config_search_path=search_path) return config_loader diff --git a/project/utils/types/__init__.py b/project/utils/typing_utils/__init__.py similarity index 100% rename from project/utils/types/__init__.py rename to project/utils/typing_utils/__init__.py diff --git a/project/utils/types/protocols.py b/project/utils/typing_utils/protocols.py similarity index 100% rename from project/utils/types/protocols.py rename to project/utils/typing_utils/protocols.py diff --git a/project/utils/utils.py b/project/utils/utils.py index 2fdd358e..95040b50 100644 --- a/project/utils/utils.py +++ b/project/utils/utils.py @@ -14,7 +14,7 @@ from omegaconf import DictConfig, OmegaConf from torchvision import transforms -from project.utils.types.protocols import ( +from project.utils.typing_utils.protocols import ( DataModule, ) diff --git a/pyproject.toml b/pyproject.toml index 3c3863aa..68c9669c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ "matplotlib>=3.9.2", "torch-jax-interop>=0.0.7", "pydantic>=2.8.2", + "simple-parsing>=0.1.5", ] readme = "README.md" requires-python = ">= 3.10" @@ -35,6 +36,7 @@ docs = [ "mkdocs-material>=9.5.31", "black>=24.8.0", "mkdocs>=1.6.0", + "mkdocs-video>=1.5.0", ] gpu = ["jax[cuda12]>=0.4.31"] @@ -47,13 +49,14 @@ build-backend = "hatchling.build" testpaths = ["project", "docs"] norecursedirs = [".venv"] # todo: look into using https://github.com/scientific-python/pytest-doctestplus +env = ["CUBLAS_WORKSPACE_CONFIG=:4096:8"] + + addopts = ["--doctest-modules", "--stats-rounding-precision=3"] markers = [ "incremental: incremental testing support", "parametrize_when_used: Parametrize a test only when the argument is used.", ] - - [tool.ruff] line-length = 99 indent-width = 4 @@ -77,6 +80,8 @@ dev-dependencies = [ "pytest-testmon>=2.1.1", "pre-commit>=3.8.0", "mktestdocs>=0.2.2", + # TODO: weird. Don't understand why this one is required! + "tomli>=2.0.1", ] [tool.hatch.metadata] @@ -90,3 +95,4 @@ serve-docs = "mkdocs serve" check_torch_gpu = "python -c 'import torch; print(torch.cuda.is_available())'" check_jax_gpu = "python -c 'import jax.numpy as jnp; print(jnp.zeros(1).devices())'" tests = "pytest -x -v" +auto_schema = "python project/utils/auto_schema.py ." diff --git a/requirements-dev.lock b/requirements-dev.lock index 878a3b15..9a952789 100644 --- a/requirements-dev.lock +++ b/requirements-dev.lock @@ -10,50 +10,50 @@ # universal: true -e file:. -absl-py==2.1.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +absl-py==2.1.0 # via chex # via optax # via orbax-checkpoint -aiohappyeyeballs==2.3.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +aiohappyeyeballs==2.3.5 # via aiohttp -aiohttp==3.10.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +aiohttp==3.10.3 # via fsspec -aiosignal==1.3.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +aiosignal==1.3.1 # via aiohttp -annotated-types==0.7.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +annotated-types==0.7.0 # via pydantic -antlr4-python3-runtime==4.9.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +antlr4-python3-runtime==4.9.3 # via hydra-core # via omegaconf -async-timeout==4.0.3 ; python_version < '3.11' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +async-timeout==4.0.3 ; python_full_version < '3.11' # via aiohttp -attrs==24.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +attrs==24.2.0 # via aiohttp -babel==2.16.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +babel==2.16.0 # via mkdocs-material -beautifulsoup4==4.12.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +beautifulsoup4==4.12.3 # via gdown -black==24.8.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +black==24.8.0 # via research-project-template -bracex==2.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +bracex==2.5 # via wcmatch -certifi==2024.7.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +certifi==2024.7.4 # via requests # via sentry-sdk -cfgv==3.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +cfgv==3.4.0 # via pre-commit -charset-normalizer==3.3.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +charset-normalizer==3.3.2 # via requests -chex==0.1.86 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +chex==0.1.86 # via optax -click==8.1.7 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +click==8.1.7 # via black # via mkdocs # via mkdocstrings # via wandb -cloudpickle==3.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +cloudpickle==3.0.0 # via submitit -colorama==0.4.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_system == 'Windows' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (sys_platform == 'win32' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +colorama==0.4.6 # via click # via colorlog # via griffe @@ -61,72 +61,74 @@ colorama==0.4.6 ; (python_version < '3.13' and sys_platform == 'linux') or (pyth # via mkdocs-material # via pytest # via tqdm -colorlog==6.8.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +colorlog==6.8.2 # via hydra-colorlog -contourpy==1.2.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +contourpy==1.2.1 # via matplotlib -coverage==7.6.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +coverage==7.6.1 # via pytest-cov # via pytest-testmon -cycler==0.12.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +cycler==0.12.1 # via matplotlib -distlib==0.3.8 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +distlib==0.3.8 # via virtualenv -docker-pycreds==0.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +docker-pycreds==0.4.0 # via wandb -etils==1.7.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +docstring-parser==0.16 + # via simple-parsing +etils==1.7.0 # via optax # via orbax-checkpoint -exceptiongroup==1.2.2 ; python_version < '3.11' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +exceptiongroup==1.2.2 ; python_full_version < '3.11' # via pytest -execnet==2.1.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +execnet==2.1.1 # via pytest-xdist -filelock==3.15.4 ; (python_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux') or (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +filelock==3.15.4 # via gdown # via torch # via triton # via virtualenv -flax==0.8.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +flax==0.8.5 # via torch-jax-interop -fonttools==4.53.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +fonttools==4.53.1 # via matplotlib -frozenlist==1.4.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +frozenlist==1.4.1 # via aiohttp # via aiosignal -fsspec==2024.6.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +fsspec==2024.6.1 # via etils # via lightning # via pytorch-lightning # via torch -gdown==5.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +gdown==5.2.0 # via research-project-template -ghp-import==2.1.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +ghp-import==2.1.0 # via mkdocs -gitdb==4.0.11 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +gitdb==4.0.11 # via gitpython -gitpython==3.1.43 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +gitpython==3.1.43 # via wandb -griffe==0.49.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +griffe==0.49.0 # via mkdocstrings-python -hydra-colorlog==1.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +hydra-colorlog==1.2.0 # via research-project-template -hydra-core==1.3.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +hydra-core==1.3.2 # via hydra-colorlog # via hydra-submitit-launcher # via hydra-zen # via research-project-template -hydra-submitit-launcher==1.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +hydra-submitit-launcher==1.2.0 # via research-project-template -hydra-zen==0.13.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +hydra-zen==0.13.0 # via research-project-template -identify==2.6.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +identify==2.6.0 # via pre-commit -idna==3.7 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +idna==3.7 # via requests # via yarl -importlib-resources==6.4.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +importlib-resources==6.4.2 # via etils -iniconfig==2.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +iniconfig==2.0.0 # via pytest jax==0.4.31 # via chex @@ -136,100 +138,105 @@ jax==0.4.31 # via pytorch2jax # via research-project-template # via torch-jax-interop -jax-cuda12-pjrt==0.4.31 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +jax-cuda12-pjrt==0.4.31 # via jax-cuda12-plugin jax-cuda12-plugin==0.4.31 # via jax -jaxlib==0.4.31 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +jaxlib==0.4.31 # via chex # via jax # via optax # via orbax-checkpoint # via pytorch2jax -jinja2==3.1.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +jinja2==3.1.4 # via mkdocs # via mkdocs-material # via mkdocstrings # via torch -kiwisolver==1.4.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +kiwisolver==1.4.5 # via matplotlib -lightning==2.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +lightning==2.4.0 # via research-project-template -lightning-utilities==0.11.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +lightning-utilities==0.11.6 # via lightning # via pytorch-lightning # via torchmetrics -markdown==3.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +lxml==5.3.0 + # via mkdocs-video +markdown==3.6 # via mkdocs # via mkdocs-autorefs # via mkdocs-material # via mkdocstrings # via pymdown-extensions -markdown-it-py==3.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +markdown-it-py==3.0.0 # via rich -markupsafe==2.1.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +markupsafe==2.1.5 # via jinja2 # via mkdocs # via mkdocs-autorefs # via mkdocstrings -matplotlib==3.9.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +matplotlib==3.9.2 # via research-project-template -mdurl==0.1.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mdurl==0.1.2 # via markdown-it-py -mergedeep==1.3.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mergedeep==1.3.4 # via mkdocs # via mkdocs-get-deps -mkdocs==1.6.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs==1.6.0 # via mkdocs-autorefs # via mkdocs-awesome-pages-plugin # via mkdocs-gen-files # via mkdocs-literate-nav # via mkdocs-material + # via mkdocs-video # via mkdocstrings # via research-project-template -mkdocs-autorefs==1.0.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-autorefs==1.0.1 # via mkdocstrings -mkdocs-awesome-pages-plugin==2.9.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-awesome-pages-plugin==2.9.3 # via research-project-template -mkdocs-gen-files==0.5.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-gen-files==0.5.0 # via research-project-template -mkdocs-get-deps==0.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-get-deps==0.2.0 # via mkdocs -mkdocs-literate-nav==0.6.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-literate-nav==0.6.1 # via research-project-template -mkdocs-material==9.5.31 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-material==9.5.31 # via research-project-template -mkdocs-material-extensions==1.3.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-material-extensions==1.3.1 # via mkdocs-material -mkdocstrings==0.25.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-video==1.5.0 + # via research-project-template +mkdocstrings==0.25.2 # via mkdocstrings-python # via research-project-template -mkdocstrings-python==1.10.8 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocstrings-python==1.10.8 # via mkdocstrings -mktestdocs==0.2.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') -ml-dtypes==0.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mktestdocs==0.2.2 +ml-dtypes==0.4.0 # via jax # via jaxlib # via tensorstore -mpmath==1.3.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mpmath==1.3.0 # via sympy -msgpack==1.0.8 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +msgpack==1.0.8 # via flax # via orbax-checkpoint -multidict==6.0.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +multidict==6.0.5 # via aiohttp # via yarl -mypy-extensions==1.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mypy-extensions==1.0.0 # via black -natsort==8.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +natsort==8.4.0 # via mkdocs-awesome-pages-plugin -nest-asyncio==1.6.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +nest-asyncio==1.6.0 # via orbax-checkpoint -networkx==3.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +networkx==3.3 # via torch -nodeenv==1.9.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +nodeenv==1.9.1 # via pre-commit -numpy==1.26.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or ((python_version >= '3.11' or (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or ((python_version >= '3.12' or (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +numpy==1.26.4 # via chex # via contourpy # via flax @@ -245,55 +252,55 @@ numpy==1.26.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python # via tensorstore # via torchmetrics # via torchvision -nvidia-cublas-cu12==12.1.3.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))))) +nvidia-cublas-cu12==12.1.3.1 # via jax-cuda12-plugin # via nvidia-cudnn-cu12 # via nvidia-cusolver-cu12 # via torch -nvidia-cuda-cupti-cu12==12.1.105 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-cuda-cupti-cu12==12.1.105 # via jax-cuda12-plugin # via torch -nvidia-cuda-nvcc-cu12==12.6.20 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +nvidia-cuda-nvcc-cu12==12.6.20 # via jax-cuda12-plugin -nvidia-cuda-nvrtc-cu12==12.1.105 ; platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +nvidia-cuda-nvrtc-cu12==12.1.105 ; platform_machine == 'x86_64' and platform_system == 'Linux' # via torch -nvidia-cuda-runtime-cu12==12.1.105 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-cuda-runtime-cu12==12.1.105 # via jax-cuda12-plugin # via torch -nvidia-cudnn-cu12==9.1.0.70 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-cudnn-cu12==9.1.0.70 # via jax-cuda12-plugin # via torch -nvidia-cufft-cu12==11.0.2.54 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-cufft-cu12==11.0.2.54 # via jax-cuda12-plugin # via torch -nvidia-curand-cu12==10.3.2.106 ; platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +nvidia-curand-cu12==10.3.2.106 ; platform_machine == 'x86_64' and platform_system == 'Linux' # via torch -nvidia-cusolver-cu12==11.4.5.107 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-cusolver-cu12==11.4.5.107 # via jax-cuda12-plugin # via torch -nvidia-cusparse-cu12==12.1.0.106 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))))) +nvidia-cusparse-cu12==12.1.0.106 # via jax-cuda12-plugin # via nvidia-cusolver-cu12 # via torch -nvidia-nccl-cu12==2.20.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-nccl-cu12==2.20.5 # via jax-cuda12-plugin # via torch -nvidia-nvjitlink-cu12==12.6.20 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))))) or (((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))))))) +nvidia-nvjitlink-cu12==12.6.20 # via jax-cuda12-plugin # via nvidia-cusolver-cu12 # via nvidia-cusparse-cu12 -nvidia-nvtx-cu12==12.1.105 ; platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +nvidia-nvtx-cu12==12.1.105 ; platform_machine == 'x86_64' and platform_system == 'Linux' # via torch -omegaconf==2.3.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +omegaconf==2.3.0 # via hydra-core # via hydra-zen -opt-einsum==3.3.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +opt-einsum==3.3.0 # via jax -optax==0.2.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +optax==0.2.3 # via flax -orbax-checkpoint==0.5.23 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +orbax-checkpoint==0.5.23 # via flax -packaging==24.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +packaging==24.1 # via black # via hydra-core # via lightning @@ -303,45 +310,45 @@ packaging==24.1 ; (python_version < '3.13' and sys_platform == 'linux') or (pyth # via pytest # via pytorch-lightning # via torchmetrics -paginate==0.5.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +paginate==0.5.6 # via mkdocs-material -pathspec==0.12.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pathspec==0.12.1 # via black # via mkdocs -pillow==10.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pillow==10.4.0 # via matplotlib # via torchvision -platformdirs==4.2.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +platformdirs==4.2.2 # via black # via mkdocs-get-deps # via mkdocstrings # via virtualenv # via wandb -pluggy==1.5.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pluggy==1.5.0 # via pytest -pre-commit==3.8.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') -protobuf==5.27.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pre-commit==3.8.0 +protobuf==5.27.3 # via orbax-checkpoint # via wandb -psutil==6.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +psutil==6.0.0 # via wandb -py-cpuinfo==9.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +py-cpuinfo==9.0.0 # via pytest-benchmark -pydantic==2.8.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pydantic==2.8.2 # via research-project-template -pydantic-core==2.20.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pydantic-core==2.20.1 # via pydantic -pygments==2.18.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pygments==2.18.0 # via mkdocs-material # via rich -pymdown-extensions==10.9 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pymdown-extensions==10.9 # via mkdocs-material # via mkdocstrings -pyparsing==3.1.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pyparsing==3.1.2 # via matplotlib -pysocks==1.7.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pysocks==1.7.1 # via requests -pytest==8.3.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pytest==8.3.2 # via pytest-benchmark # via pytest-cov # via pytest-datadir @@ -351,25 +358,25 @@ pytest==8.3.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python # via pytest-testmon # via pytest-timeout # via pytest-xdist -pytest-benchmark==4.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') -pytest-cov==5.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') -pytest-datadir==1.5.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pytest-benchmark==4.0.0 +pytest-cov==5.0.0 +pytest-datadir==1.5.0 # via pytest-regressions -pytest-env==1.1.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') -pytest-regressions==2.5.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pytest-env==1.1.3 +pytest-regressions==2.5.0 # via tensor-regression -pytest-skip-slow==0.0.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') -pytest-testmon==2.1.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') -pytest-timeout==2.3.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') -pytest-xdist==3.6.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') -python-dateutil==2.9.0.post0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pytest-skip-slow==0.0.5 +pytest-testmon==2.1.1 +pytest-timeout==2.3.1 +pytest-xdist==3.6.1 +python-dateutil==2.9.0.post0 # via ghp-import # via matplotlib -pytorch-lightning==2.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pytorch-lightning==2.4.0 # via lightning -pytorch2jax==0.1.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pytorch2jax==0.1.0 # via torch-jax-interop -pyyaml==6.0.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pyyaml==6.0.2 # via flax # via lightning # via mkdocs @@ -382,52 +389,54 @@ pyyaml==6.0.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python # via pytorch-lightning # via pyyaml-env-tag # via wandb -pyyaml-env-tag==0.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pyyaml-env-tag==0.1 # via mkdocs -regex==2024.7.24 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +regex==2024.7.24 # via mkdocs-material -requests==2.32.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +requests==2.32.3 # via gdown # via mkdocs-material # via wandb -rich==13.7.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +rich==13.7.1 # via flax # via research-project-template -ruff==0.6.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') -scipy==1.14.0 ; (python_version >= '3.12' or (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +ruff==0.6.0 +scipy==1.14.0 # via jax # via jaxlib -sentry-sdk==2.13.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +sentry-sdk==2.13.0 # via wandb -setproctitle==1.3.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +setproctitle==1.3.3 # via wandb -setuptools==72.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.12' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +setuptools==72.2.0 # via chex # via lightning-utilities # via wandb -six==1.16.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +simple-parsing==0.1.5 + # via research-project-template +six==1.16.0 # via docker-pycreds # via python-dateutil -smmap==5.0.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +smmap==5.0.1 # via gitdb -soupsieve==2.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +soupsieve==2.6 # via beautifulsoup4 -submitit==1.5.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +submitit==1.5.1 # via hydra-submitit-launcher -sympy==1.13.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +sympy==1.13.2 # via torch -tensor-regression==0.0.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') -tensorstore==0.1.64 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +tensor-regression==0.0.6 +tensorstore==0.1.64 # via flax # via orbax-checkpoint -tomli==2.0.1 ; (python_full_version <= '3.11.0a6' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (python_version < '3.11' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +tomli==2.0.1 # via black # via coverage # via pytest # via pytest-env -toolz==0.12.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +toolz==0.12.1 # via chex -torch==2.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +torch==2.4.0 # via lightning # via pytorch-lightning # via pytorch2jax @@ -436,21 +445,21 @@ torch==2.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_ # via torch-jax-interop # via torchmetrics # via torchvision -torch-jax-interop==0.0.7 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +torch-jax-interop==0.0.7 # via research-project-template -torchmetrics==1.4.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +torchmetrics==1.4.1 # via lightning # via pytorch-lightning -torchvision==0.19.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +torchvision==0.19.0 # via research-project-template -tqdm==4.66.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +tqdm==4.66.5 # via gdown # via lightning # via pytorch-lightning # via research-project-template -triton==3.0.0 ; python_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux' +triton==3.0.0 ; python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux' # via torch -typing-extensions==4.12.2 ; (python_version < '3.11' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +typing-extensions==4.12.2 # via black # via chex # via etils @@ -462,20 +471,21 @@ typing-extensions==4.12.2 ; (python_version < '3.11' and ((python_version < '3.1 # via pydantic # via pydantic-core # via pytorch-lightning + # via simple-parsing # via submitit # via torch -urllib3==2.2.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +urllib3==2.2.2 # via requests # via sentry-sdk -virtualenv==20.26.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +virtualenv==20.26.3 # via pre-commit -wandb==0.17.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +wandb==0.17.6 # via research-project-template -watchdog==4.0.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +watchdog==4.0.2 # via mkdocs -wcmatch==9.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +wcmatch==9.0 # via mkdocs-awesome-pages-plugin -yarl==1.9.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +yarl==1.9.4 # via aiohttp -zipp==3.20.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +zipp==3.20.0 # via etils diff --git a/requirements.lock b/requirements.lock index 1bf0a5cd..5a034d4f 100644 --- a/requirements.lock +++ b/requirements.lock @@ -10,106 +10,108 @@ # universal: true -e file:. -absl-py==2.1.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +absl-py==2.1.0 # via chex # via optax # via orbax-checkpoint -aiohappyeyeballs==2.3.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +aiohappyeyeballs==2.3.5 # via aiohttp -aiohttp==3.10.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +aiohttp==3.10.3 # via fsspec -aiosignal==1.3.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +aiosignal==1.3.1 # via aiohttp -annotated-types==0.7.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +annotated-types==0.7.0 # via pydantic -antlr4-python3-runtime==4.9.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +antlr4-python3-runtime==4.9.3 # via hydra-core # via omegaconf -async-timeout==4.0.3 ; python_version < '3.11' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +async-timeout==4.0.3 ; python_full_version < '3.11' # via aiohttp -attrs==24.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +attrs==24.2.0 # via aiohttp -babel==2.16.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +babel==2.16.0 # via mkdocs-material -beautifulsoup4==4.12.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +beautifulsoup4==4.12.3 # via gdown -black==24.8.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +black==24.8.0 # via research-project-template -bracex==2.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +bracex==2.5 # via wcmatch -certifi==2024.7.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +certifi==2024.7.4 # via requests # via sentry-sdk -charset-normalizer==3.3.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +charset-normalizer==3.3.2 # via requests -chex==0.1.86 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +chex==0.1.86 # via optax -click==8.1.7 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +click==8.1.7 # via black # via mkdocs # via mkdocstrings # via wandb -cloudpickle==3.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +cloudpickle==3.0.0 # via submitit -colorama==0.4.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_system == 'Windows' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (sys_platform == 'win32' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +colorama==0.4.6 # via click # via colorlog # via griffe # via mkdocs # via mkdocs-material # via tqdm -colorlog==6.8.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +colorlog==6.8.2 # via hydra-colorlog -contourpy==1.2.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +contourpy==1.2.1 # via matplotlib -cycler==0.12.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +cycler==0.12.1 # via matplotlib -docker-pycreds==0.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +docker-pycreds==0.4.0 # via wandb -etils==1.7.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +docstring-parser==0.16 + # via simple-parsing +etils==1.7.0 # via optax # via orbax-checkpoint -filelock==3.15.4 ; (python_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux') or (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +filelock==3.15.4 # via gdown # via torch # via triton -flax==0.8.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +flax==0.8.5 # via torch-jax-interop -fonttools==4.53.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +fonttools==4.53.1 # via matplotlib -frozenlist==1.4.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +frozenlist==1.4.1 # via aiohttp # via aiosignal -fsspec==2024.6.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +fsspec==2024.6.1 # via etils # via lightning # via pytorch-lightning # via torch -gdown==5.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +gdown==5.2.0 # via research-project-template -ghp-import==2.1.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +ghp-import==2.1.0 # via mkdocs -gitdb==4.0.11 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +gitdb==4.0.11 # via gitpython -gitpython==3.1.43 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +gitpython==3.1.43 # via wandb -griffe==0.49.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +griffe==0.49.0 # via mkdocstrings-python -hydra-colorlog==1.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +hydra-colorlog==1.2.0 # via research-project-template -hydra-core==1.3.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +hydra-core==1.3.2 # via hydra-colorlog # via hydra-submitit-launcher # via hydra-zen # via research-project-template -hydra-submitit-launcher==1.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +hydra-submitit-launcher==1.2.0 # via research-project-template -hydra-zen==0.13.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +hydra-zen==0.13.0 # via research-project-template -idna==3.7 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +idna==3.7 # via requests # via yarl -importlib-resources==6.4.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +importlib-resources==6.4.2 # via etils jax==0.4.31 # via chex @@ -119,97 +121,102 @@ jax==0.4.31 # via pytorch2jax # via research-project-template # via torch-jax-interop -jax-cuda12-pjrt==0.4.31 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +jax-cuda12-pjrt==0.4.31 # via jax-cuda12-plugin jax-cuda12-plugin==0.4.31 # via jax -jaxlib==0.4.31 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +jaxlib==0.4.31 # via chex # via jax # via optax # via orbax-checkpoint # via pytorch2jax -jinja2==3.1.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +jinja2==3.1.4 # via mkdocs # via mkdocs-material # via mkdocstrings # via torch -kiwisolver==1.4.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +kiwisolver==1.4.5 # via matplotlib -lightning==2.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +lightning==2.4.0 # via research-project-template -lightning-utilities==0.11.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +lightning-utilities==0.11.6 # via lightning # via pytorch-lightning # via torchmetrics -markdown==3.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +lxml==5.3.0 + # via mkdocs-video +markdown==3.6 # via mkdocs # via mkdocs-autorefs # via mkdocs-material # via mkdocstrings # via pymdown-extensions -markdown-it-py==3.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +markdown-it-py==3.0.0 # via rich -markupsafe==2.1.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +markupsafe==2.1.5 # via jinja2 # via mkdocs # via mkdocs-autorefs # via mkdocstrings -matplotlib==3.9.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +matplotlib==3.9.2 # via research-project-template -mdurl==0.1.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mdurl==0.1.2 # via markdown-it-py -mergedeep==1.3.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mergedeep==1.3.4 # via mkdocs # via mkdocs-get-deps -mkdocs==1.6.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs==1.6.0 # via mkdocs-autorefs # via mkdocs-awesome-pages-plugin # via mkdocs-gen-files # via mkdocs-literate-nav # via mkdocs-material + # via mkdocs-video # via mkdocstrings # via research-project-template -mkdocs-autorefs==1.0.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-autorefs==1.0.1 # via mkdocstrings -mkdocs-awesome-pages-plugin==2.9.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-awesome-pages-plugin==2.9.3 # via research-project-template -mkdocs-gen-files==0.5.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-gen-files==0.5.0 # via research-project-template -mkdocs-get-deps==0.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-get-deps==0.2.0 # via mkdocs -mkdocs-literate-nav==0.6.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-literate-nav==0.6.1 # via research-project-template -mkdocs-material==9.5.31 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-material==9.5.31 # via research-project-template -mkdocs-material-extensions==1.3.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-material-extensions==1.3.1 # via mkdocs-material -mkdocstrings==0.25.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocs-video==1.5.0 + # via research-project-template +mkdocstrings==0.25.2 # via mkdocstrings-python # via research-project-template -mkdocstrings-python==1.10.8 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mkdocstrings-python==1.10.8 # via mkdocstrings -ml-dtypes==0.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +ml-dtypes==0.4.0 # via jax # via jaxlib # via tensorstore -mpmath==1.3.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mpmath==1.3.0 # via sympy -msgpack==1.0.8 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +msgpack==1.0.8 # via flax # via orbax-checkpoint -multidict==6.0.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +multidict==6.0.5 # via aiohttp # via yarl -mypy-extensions==1.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +mypy-extensions==1.0.0 # via black -natsort==8.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +natsort==8.4.0 # via mkdocs-awesome-pages-plugin -nest-asyncio==1.6.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +nest-asyncio==1.6.0 # via orbax-checkpoint -networkx==3.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +networkx==3.3 # via torch -numpy==1.26.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or ((python_version >= '3.11' or (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or ((python_version >= '3.12' or (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +numpy==1.26.4 # via chex # via contourpy # via flax @@ -224,55 +231,55 @@ numpy==1.26.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python # via tensorstore # via torchmetrics # via torchvision -nvidia-cublas-cu12==12.1.3.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))))) +nvidia-cublas-cu12==12.1.3.1 # via jax-cuda12-plugin # via nvidia-cudnn-cu12 # via nvidia-cusolver-cu12 # via torch -nvidia-cuda-cupti-cu12==12.1.105 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-cuda-cupti-cu12==12.1.105 # via jax-cuda12-plugin # via torch -nvidia-cuda-nvcc-cu12==12.6.20 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +nvidia-cuda-nvcc-cu12==12.6.20 # via jax-cuda12-plugin -nvidia-cuda-nvrtc-cu12==12.1.105 ; platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +nvidia-cuda-nvrtc-cu12==12.1.105 ; platform_machine == 'x86_64' and platform_system == 'Linux' # via torch -nvidia-cuda-runtime-cu12==12.1.105 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-cuda-runtime-cu12==12.1.105 # via jax-cuda12-plugin # via torch -nvidia-cudnn-cu12==9.1.0.70 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-cudnn-cu12==9.1.0.70 # via jax-cuda12-plugin # via torch -nvidia-cufft-cu12==11.0.2.54 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-cufft-cu12==11.0.2.54 # via jax-cuda12-plugin # via torch -nvidia-curand-cu12==10.3.2.106 ; platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +nvidia-curand-cu12==10.3.2.106 ; platform_machine == 'x86_64' and platform_system == 'Linux' # via torch -nvidia-cusolver-cu12==11.4.5.107 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-cusolver-cu12==11.4.5.107 # via jax-cuda12-plugin # via torch -nvidia-cusparse-cu12==12.1.0.106 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))))) +nvidia-cusparse-cu12==12.1.0.106 # via jax-cuda12-plugin # via nvidia-cusolver-cu12 # via torch -nvidia-nccl-cu12==2.20.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) +nvidia-nccl-cu12==2.20.5 # via jax-cuda12-plugin # via torch -nvidia-nvjitlink-cu12==12.6.20 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))))) or (((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') or (platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))))))) +nvidia-nvjitlink-cu12==12.6.20 # via jax-cuda12-plugin # via nvidia-cusolver-cu12 # via nvidia-cusparse-cu12 -nvidia-nvtx-cu12==12.1.105 ; platform_machine == 'x86_64' and platform_system == 'Linux' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +nvidia-nvtx-cu12==12.1.105 ; platform_machine == 'x86_64' and platform_system == 'Linux' # via torch -omegaconf==2.3.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +omegaconf==2.3.0 # via hydra-core # via hydra-zen -opt-einsum==3.3.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +opt-einsum==3.3.0 # via jax -optax==0.2.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +optax==0.2.3 # via flax -orbax-checkpoint==0.5.23 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +orbax-checkpoint==0.5.23 # via flax -packaging==24.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +packaging==24.1 # via black # via hydra-core # via lightning @@ -281,46 +288,46 @@ packaging==24.1 ; (python_version < '3.13' and sys_platform == 'linux') or (pyth # via mkdocs # via pytorch-lightning # via torchmetrics -paginate==0.5.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +paginate==0.5.6 # via mkdocs-material -pathspec==0.12.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pathspec==0.12.1 # via black # via mkdocs -pillow==10.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pillow==10.4.0 # via matplotlib # via torchvision -platformdirs==4.2.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +platformdirs==4.2.2 # via black # via mkdocs-get-deps # via mkdocstrings # via wandb -protobuf==5.27.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +protobuf==5.27.3 # via orbax-checkpoint # via wandb -psutil==6.0.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +psutil==6.0.0 # via wandb -pydantic==2.8.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pydantic==2.8.2 # via research-project-template -pydantic-core==2.20.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pydantic-core==2.20.1 # via pydantic -pygments==2.18.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pygments==2.18.0 # via mkdocs-material # via rich -pymdown-extensions==10.9 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pymdown-extensions==10.9 # via mkdocs-material # via mkdocstrings -pyparsing==3.1.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pyparsing==3.1.2 # via matplotlib -pysocks==1.7.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pysocks==1.7.1 # via requests -python-dateutil==2.9.0.post0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +python-dateutil==2.9.0.post0 # via ghp-import # via matplotlib -pytorch-lightning==2.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pytorch-lightning==2.4.0 # via lightning -pytorch2jax==0.1.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pytorch2jax==0.1.0 # via torch-jax-interop -pyyaml==6.0.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pyyaml==6.0.2 # via flax # via lightning # via mkdocs @@ -331,47 +338,49 @@ pyyaml==6.0.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python # via pytorch-lightning # via pyyaml-env-tag # via wandb -pyyaml-env-tag==0.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +pyyaml-env-tag==0.1 # via mkdocs -regex==2024.7.24 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +regex==2024.7.24 # via mkdocs-material -requests==2.32.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +requests==2.32.3 # via gdown # via mkdocs-material # via wandb -rich==13.7.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +rich==13.7.1 # via flax # via research-project-template -scipy==1.14.0 ; (python_version >= '3.12' or (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +scipy==1.14.0 # via jax # via jaxlib -sentry-sdk==2.13.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +sentry-sdk==2.13.0 # via wandb -setproctitle==1.3.3 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +setproctitle==1.3.3 # via wandb -setuptools==72.2.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.12' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +setuptools==72.2.0 # via chex # via lightning-utilities # via wandb -six==1.16.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +simple-parsing==0.1.5 + # via research-project-template +six==1.16.0 # via docker-pycreds # via python-dateutil -smmap==5.0.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +smmap==5.0.1 # via gitdb -soupsieve==2.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +soupsieve==2.6 # via beautifulsoup4 -submitit==1.5.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +submitit==1.5.1 # via hydra-submitit-launcher -sympy==1.13.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +sympy==1.13.2 # via torch -tensorstore==0.1.64 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +tensorstore==0.1.64 # via flax # via orbax-checkpoint -tomli==2.0.1 ; python_version < '3.11' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux')) +tomli==2.0.1 ; python_full_version < '3.11' # via black -toolz==0.12.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +toolz==0.12.1 # via chex -torch==2.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +torch==2.4.0 # via lightning # via pytorch-lightning # via pytorch2jax @@ -379,21 +388,21 @@ torch==2.4.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_ # via torch-jax-interop # via torchmetrics # via torchvision -torch-jax-interop==0.0.7 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +torch-jax-interop==0.0.7 # via research-project-template -torchmetrics==1.4.1 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +torchmetrics==1.4.1 # via lightning # via pytorch-lightning -torchvision==0.19.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +torchvision==0.19.0 # via research-project-template -tqdm==4.66.5 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +tqdm==4.66.5 # via gdown # via lightning # via pytorch-lightning # via research-project-template -triton==3.0.0 ; python_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux' +triton==3.0.0 ; python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux' # via torch -typing-extensions==4.12.2 ; (python_version < '3.11' and ((python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux'))) or (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +typing-extensions==4.12.2 # via black # via chex # via etils @@ -405,18 +414,19 @@ typing-extensions==4.12.2 ; (python_version < '3.11' and ((python_version < '3.1 # via pydantic # via pydantic-core # via pytorch-lightning + # via simple-parsing # via submitit # via torch -urllib3==2.2.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +urllib3==2.2.2 # via requests # via sentry-sdk -wandb==0.17.6 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +wandb==0.17.6 # via research-project-template -watchdog==4.0.2 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +watchdog==4.0.2 # via mkdocs -wcmatch==9.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +wcmatch==9.0 # via mkdocs-awesome-pages-plugin -yarl==1.9.4 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +yarl==1.9.4 # via aiohttp -zipp==3.20.0 ; (python_version < '3.13' and sys_platform == 'linux') or (python_version < '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version < '3.13' and sys_platform != 'linux') or (python_version >= '3.13' and sys_platform == 'linux') or (python_version >= '3.13' and sys_platform == 'linux' and (python_version <= '3.9' or sys_platform != 'linux')) or (python_version >= '3.13' and sys_platform != 'linux') +zipp==3.20.0 # via etils