lebrice

from __future__ import annotations

class Value:
    """ stores a single scalar value and its gradient """

    def __init__(self, data, _parents: tuple[Value, ...]=(), _op=''):
        self.data = data
        self.grad = 0
        # internal variables used for autograd graph construction
        self._backward = lambda: None

lebrice

from __future__ import annotations

import copy
import functools
import math
from collections import OrderedDict
from typing import Sequence

import torch
from torch import Tensor, nn

lebrice

"""
"""
from __future__ import annotations

import datetime
import itertools

# NOTE: Need to import cv2 to prevent a loading error for GLIBCXX with ffcv.
import cv2  # noqa

lebrice

""" ImageNet datamodule that uses FFCV. """

from __future__ import annotations

import typing
from collections.abc import Iterable, Sequence
from pathlib import Path
from typing import Any, Callable, TypeVar

import cv2  # noqa

lebrice

from __future__ import annotations
import itertools
from pytorch_lightning import Trainer
import torch
import numpy as np
from pytorch_lightning import LightningModule
from torch import nn, Tensor
import pytest
from .imagenet_ffcv import ImagenetFfcvDataModule

lebrice

from __future__ import annotations
from pytorch_lightning import Trainer
from typing_extensions import Unpack, ParamSpec
from typing import Callable, TypedDict

# Option A: TypedDict
# --> lots of code duplication!


class TrainerConfig(TypedDict, total=False):

lebrice

""" ImageNet datamodule adapted to the Mila cluster.

Can be used either with a PyTorch-Lightning Trainer, or by itself to easily get efficient
dataloaders for the ImageNet dataset.

Requirements (these are the versions I'm using, but this can probably be loosened a bit).
- pytorch-lightning==1.6.0
- lightning-bolts==0.5
"""

lebrice

"""Sets up a user cache directory for commonly used libraries, while reusing shared cache entries.

Use this to avoid having to download files to the $HOME directory, as well as to remove
duplicated downloads and free up space in your $HOME and $SCRATCH directories.

The user cache directory should be writeable, and doesn't need to be empty.
This command adds symlinks to (some of) the files contained in the *shared* cache directory to this
user cache directory.

The shared cache directory should be readable (e.g. a directory containing frequently-downloaded

lebrice

from __future__ import annotations

# Context: Dataset is on GPU memory.
from typing import Iterable
import torch
from torch import Tensor
from torchvision.datasets import MNIST
from torch.utils.data import TensorDataset, DataLoader, Dataset, ConcatDataset

lebrice

from dataclasses import dataclass
import os
from simple_parsing import ArgumentParser
from itertools import product


@dataclass
class ProblemConfig:
    dataset: int = 0  # Which dataset ID to use.
    rank: int = 0  # The rank of some matrix