dvruette

from contextlib import contextmanager

@contextmanager
def set_numpy(name):
    if name == "numpy":
        import numpy as np
        yield np
    elif name == "jax":
        import jax.numpy as jnp
        yield jnp

dvruette

diff --git a/grid.py b/grid.py
index f9f1557..8eafb91 100755
--- a/grid.py
+++ b/grid.py
@@ -5,8 +5,10 @@
 
 # Written by Francois Fleuret <[email protected]>
 
-import math
-import torch, torchvision

dvruette

import copy
import torch
import torch.nn as nn


class DecayToInit(nn.Module):
    def __init__(self, param: torch.Tensor):
        super().__init__()
        self.register_buffer("param", param)

dvruette

"""
Based on: https://github.com/lucidrains/lion-pytorch/blob/main/lion_pytorch/lion_pytorch.py
"""

from typing import Tuple, Optional, Callable

import torch
from torch.optim.optimizer import Optimizer