Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Resolve conflict with PyTorch >= 2.2 #1

Closed
wants to merge 2 commits into from
Closed

Resolve conflict with PyTorch >= 2.2 #1

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
48ec5fe
Resolve conflict with PyTorch >= 2.2
HuFY-dev Mar 29, 2024
5f9997f
Resolved import error with PyTorch >= 2.2
HuFY-dev Mar 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions sparse_autoencoder/optimizer/adam_with_reset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,20 @@

This reset method is useful when resampling dead neurons during training.
"""
from collections.abc import Iterator
from collections.abc import Iterable, Iterator
from typing import Any
from typing_extensions import TypeAlias

from jaxtyping import Float, Int
from torch import Tensor
from torch.nn.parameter import Parameter
from torch.optim import Adam
from torch.optim.optimizer import params_t

from sparse_autoencoder.tensor_types import Axis

# params_t was renamed to ParamsT in PyTorch 2.2, which caused import errors

Check failure on line 16 in sparse_autoencoder/optimizer/adam_with_reset.py

View workflow job for this annotation

GitHub Actions / Checks (3.11)

Ruff (I001) 

sparse_autoencoder/optimizer/adam_with_reset.py:5:1: I001 Import block is un-sorted or un-formatted
# Copied from PyTorch 2.2 with modifications for better style
ParamsT: TypeAlias = Iterable[Tensor] | Iterable[dict[str, Any]]

class AdamWithReset(Adam):
"""Adam Optimizer with a reset method.
Expand All @@ -35,7 +39,7 @@

def __init__( # (extending existing implementation)
self,
params: params_t,
params: ParamsT,
lr: float | Float[Tensor, Axis.names(Axis.SINGLE_ITEM)] = 1e-3,
betas: tuple[float, float] = (0.9, 0.999),
eps: float = 1e-8,
Expand Down
Loading