heavyball 1.1.2__tar.gz → 1.1.3__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {heavyball-1.1.2 → heavyball-1.1.3}/PKG-INFO +1 -1
- {heavyball-1.1.2 → heavyball-1.1.3}/heavyball/chainable.py +3 -4
- {heavyball-1.1.2 → heavyball-1.1.3}/heavyball.egg-info/PKG-INFO +1 -1
- {heavyball-1.1.2 → heavyball-1.1.3}/setup.py +1 -1
- {heavyball-1.1.2 → heavyball-1.1.3}/LICENSE +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/README.md +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/heavyball/__init__.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/heavyball/utils.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/heavyball.egg-info/SOURCES.txt +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/heavyball.egg-info/dependency_links.txt +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/heavyball.egg-info/requires.txt +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/heavyball.egg-info/top_level.txt +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/setup.cfg +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_bf16_params.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_bf16_q.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_bf16_storage.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_caution.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_channels_last.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_closure.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_ema.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_foreach.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_hook.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_mars.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_memory.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_merge.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_no_grad.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_psgd.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_soap.py +0 -0
- {heavyball-1.1.2 → heavyball-1.1.3}/test/test_stochastic_updates.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
import functools
|
|
2
2
|
import random
|
|
3
|
-
from typing import Optional, Union
|
|
3
|
+
from typing import Optional, Union, Literal
|
|
4
4
|
|
|
5
5
|
import torch
|
|
6
6
|
|
|
@@ -51,8 +51,7 @@ class FunctionTransform:
|
|
|
51
51
|
|
|
52
52
|
|
|
53
53
|
def _zero_guard(state, key, ref, dtype):
|
|
54
|
-
return _guard_in_state(state, key,
|
|
55
|
-
lambda: torch.zeros_like(ref, dtype=torch.float32, memory_format=torch.preserve_format))
|
|
54
|
+
return _guard_in_state(state, key, lambda: torch.zeros_like(ref, dtype=dtype, memory_format=torch.preserve_format))
|
|
56
55
|
|
|
57
56
|
|
|
58
57
|
def _storage_dtype(group):
|
|
@@ -442,7 +441,7 @@ class ChainOpt(utils.StatefulOptimizer):
|
|
|
442
441
|
|
|
443
442
|
|
|
444
443
|
use_default = object()
|
|
445
|
-
str_or_fn = Union[str, callable, None, use_default]
|
|
444
|
+
str_or_fn = Union[str, callable, None, Literal[use_default]]
|
|
446
445
|
|
|
447
446
|
|
|
448
447
|
def _get_clip_fn(name: str_or_fn, default_val: str_or_fn):
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|