This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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 | |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 <francois@fleuret.org> | |
-import math | |
-import torch, torchvision |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |