heavyball 2.1.2__tar.gz → 2.1.3__tar.gz

{heavyball-2.1.2 → heavyball-2.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: heavyball
-Version: 2.1.2
+Version: 2.1.3
 Summary: Efficient Optimizers
 Author-email: HeavyBall Authors <github.heavyball@nestler.sh>
 Project-URL: source, https://github.com/HomebrewML/HeavyBall
@@ -21,6 +21,7 @@ Requires-Dist: numpy<2.0.0
 Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
+Requires-Dist: hypothesis; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: matplotlib; extra == "dev"
 Requires-Dist: seaborn; extra == "dev"

{heavyball-2.1.2 → heavyball-2.1.3}/heavyball/helpers.py

@@ -3,7 +3,8 @@ from __future__ import annotations
 import functools
 import math
 import threading
-from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, Generator, Iterable, List, Optional, Sequence, Tuple, Union
 
 import numpy
 import numpy as np
@@ -11,7 +12,6 @@ import optuna
 import optunahub
 import pandas as pd
 import torch
-from botorch.utils.sampling import manual_seed
 from hebo.design_space.design_space import DesignSpace
 from hebo.optimizers.hebo import HEBO
 from optuna._transform import _SearchSpaceTransform
@@ -21,13 +21,6 @@ from optuna.samplers._lazy_random_state import LazyRandomState
 from optuna.study import Study
 from optuna.study._study_direction import StudyDirection
 from optuna.trial import FrozenTrial, TrialState
-from optuna_integration.botorch import (
-    ehvi_candidates_func,
-    logei_candidates_func,
-    qehvi_candidates_func,
-    qei_candidates_func,
-    qparego_candidates_func,
-)
 from torch import Tensor
 from torch.nn import functional as F
 
@@ -37,6 +30,33 @@ _MAXINT32 = (1 << 31) - 1
 _SAMPLER_KEY = "auto:sampler"
 
 
+@contextmanager
+def manual_seed(seed: int | None = None) -> Generator[None, None, None]:
+    r"""
+    Contextmanager for manual setting the torch.random seed.
+
+    Args:
+        seed: The seed to set the random number generator to.
+
+    Returns:
+        Generator
+
+    Example:
+        >>> with manual_seed(1234):
+        >>>     X = torch.rand(3)
+
+    copied as-is from https://github.com/meta-pytorch/botorch/blob/a42cd65f9b704cdb6f2ee64db99a022eb15295d5/botorch/utils/sampling.py#L53C1-L75C50 under the MIT License
+    """
+    old_state = torch.random.get_rng_state()
+    try:
+        if seed is not None:
+            torch.random.manual_seed(seed)
+        yield
+    finally:
+        if seed is not None:
+            torch.random.set_rng_state(old_state)
+
+
 class SimpleAPIBaseSampler(BaseSampler):
     def __init__(
         self,
@@ -65,6 +85,16 @@ def _get_default_candidates_func(
     """
     The original is available at https://github.com/optuna/optuna-integration/blob/156a8bc081322791015d2beefff9373ed7b24047/optuna_integration/botorch/botorch.py under the MIT License
     """
+
+    # lazy import
+    from optuna_integration.botorch import (
+        ehvi_candidates_func,
+        logei_candidates_func,
+        qehvi_candidates_func,
+        qei_candidates_func,
+        qparego_candidates_func,
+    )
+
     if n_objectives > 3 and not has_constraint and not consider_running_trials:
         return ehvi_candidates_func
     elif n_objectives > 3:

{heavyball-2.1.2 → heavyball-2.1.3}/heavyball/utils.py

@@ -47,7 +47,7 @@ _cudnn_double_backward_pattern = re.compile(
 )
 _torch_compile_double_backward_pattern = re.compile(r"compile.*does not currently support double backward")
 _fd_error = (
-    "You can accelerate startup by globally enabling finite_differences first "  #
+    "You can accelerate startup by globally enabling finite_differences first "
     "(via opt.finite_differences=True or by subclassing it)\n"
     "Original Error: "
 )
@@ -418,9 +418,13 @@ def zeropower_via_newtonschulz5(G, steps=5, eps=1e-7):
 
 
 ###### START
-# Taken from https://github.com/thinking-machines-lab/manifolds/blob/89dcae50f01af59f1e0570289474da3a2ecaa60b/src/msign.py#L47
+# Based on https://arxiv.org/pdf/2505.16932v3
+# and https://github.com/NoahAmsel/PolarExpress/blob/5454910920ca8c65afda28820cdf9e49b9436ed0/polar_express.py#L69-L82
+# and https://github.com/thinking-machines-lab/manifolds/blob/89dcae50f01af59f1e0570289474da3a2ecaa60b/src/msign.py#L47
+#
 # under the MIT License
 
+# Coefficients are from https://arxiv.org/pdf/2505.16932v3
 ABC_LIST: list[tuple[float, float, float]] = [
     (8.28721201814563, -23.595886519098837, 17.300387312530933),
     (4.107059111542203, -2.9478499167379106, 0.5448431082926601),
@@ -438,7 +442,7 @@ ABC_LIST_STABLE: list[tuple[float, float, float]] = [
 ] + [ABC_LIST[-1]]
 
 
-def msign(G: torch.Tensor, steps: int = 10) -> torch.Tensor:
+def msign(G: torch.Tensor, steps: int = 10, eps: float = 1e-7) -> torch.Tensor:
     """
     Polar Express algorithm for the matrix sign function:
     https://arxiv.org/abs/2505.16932
@@ -450,7 +454,9 @@ def msign(G: torch.Tensor, steps: int = 10) -> torch.Tensor:
     if should_transpose:
         x = x.mT
 
-    x /= x.norm(dim=(-2, -1), keepdim=True) * 1.01
+    # x = x / (x.norm(dim=(-2, -1), keepdim=True) * 1.01 + eps)
+    stochastic_divide_with_eps_(x, x.norm(dim=(-2, -1)) * 1.01, eps)
+
     for step in range(steps):
         a, b, c = ABC_LIST_STABLE[step] if step < len(ABC_LIST_STABLE) else ABC_LIST_STABLE[-1]
         s = x @ x.mT
@@ -464,7 +470,6 @@ def msign(G: torch.Tensor, steps: int = 10) -> torch.Tensor:
 
     if should_transpose:
         x = x.mT
-    x = torch.nan_to_num(x)
     return x.float()
 
 
@@ -1585,7 +1590,7 @@ def _compilable_copy_stochastic_(target: Tensor, source: Tensor):
 
 def copy_stochastic_(target: Tensor, source: Tensor):
     if target.dtype == torch.bfloat16 and source.dtype in (torch.float16, torch.float32, torch.float64):
-        _compilable_copy_stochastic_(target, source.float())
+        source = stochastic_round_(target, source)
     set_(target, source)
 
 
@@ -2412,10 +2417,11 @@ def bf16_matmul(x: Tensor, y: Tensor):
 def if_iscompiling(fn):
     base = getattr(torch, fn.__name__, None)
 
-    def _fn(x):
-        if torch.compiler.is_compiling() and hasattr(torch, fn.__name__):
-            return base(x)
-        return fn(x)
+    @functools.wraps(fn)
+    def _fn(*args, **kwargs):
+        if torch.compiler.is_compiling() and base is not None:
+            return base(*args, **kwargs)
+        return fn(*args, **kwargs)
 
     return _fn
 

{heavyball-2.1.2 → heavyball-2.1.3}/heavyball.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: heavyball
-Version: 2.1.2
+Version: 2.1.3
 Summary: Efficient Optimizers
 Author-email: HeavyBall Authors <github.heavyball@nestler.sh>
 Project-URL: source, https://github.com/HomebrewML/HeavyBall
@@ -21,6 +21,7 @@ Requires-Dist: numpy<2.0.0
 Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
+Requires-Dist: hypothesis; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: matplotlib; extra == "dev"
 Requires-Dist: seaborn; extra == "dev"

{heavyball-2.1.2 → heavyball-2.1.3}/heavyball.egg-info/SOURCES.txt

@@ -14,11 +14,14 @@ test/test_bf16_params.py
 test/test_bf16_q.py
 test/test_bf16_storage.py
 test/test_caution.py
+test/test_chainable_cpu.py
 test/test_channels_last.py
 test/test_clip.py
 test/test_closure.py
+test/test_cpu_features.py
 test/test_ema.py
 test/test_foreach.py
+test/test_helpers_cpu.py
 test/test_hook.py
 test/test_mars.py
 test/test_memory.py
@@ -27,9 +30,12 @@ test/test_merge.py
 test/test_migrate_cli.py
 test/test_nd_param.py
 test/test_no_grad.py
+test/test_optimizer_cpu_smoke.py
 test/test_psgd_precond_init_stability.py
 test/test_save_restore.py
 test/test_singular_values.py
 test/test_soap.py
 test/test_stochastic_updates.py
-test/test_toy_training.py
+test/test_toy_training.py
+test/test_utils_cpu.py
+test/test_utils_property.py

{heavyball-2.1.2 → heavyball-2.1.3}/heavyball.egg-info/requires.txt

@@ -5,6 +5,7 @@ numpy<2.0.0
 [dev]
 pre-commit
 pytest
+hypothesis
 ruff
 matplotlib
 seaborn

{heavyball-2.1.2 → heavyball-2.1.3}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "heavyball"
 description = "Efficient Optimizers"
-version = "2.1.2"
+version = "2.1.3"
 authors = [{ name = "HeavyBall Authors", email = "github.heavyball@nestler.sh" }]
 classifiers = ["Intended Audience :: Developers",
     "Intended Audience :: Science/Research",
@@ -28,7 +28,7 @@ readme = "README.md"
 requires-python = ">=3.9"
 
 [project.optional-dependencies]
-dev = ["pre-commit", "pytest", "ruff", "matplotlib", "seaborn", "pandas", "typer", "optuna", "optunahub", "hebo", "lightbench"]
+dev = ["pre-commit", "pytest", "hypothesis", "ruff", "matplotlib", "seaborn", "pandas", "typer", "optuna", "optunahub", "hebo", "lightbench"]
 
 [project.urls]
 source = "https://github.com/HomebrewML/HeavyBall"

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_bf16_params.py

@@ -8,7 +8,6 @@ from torch import nn
 from torch._dynamo import config
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 os.environ["TORCH_LOGS"] = "+recompiles"

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_bf16_q.py

@@ -5,7 +5,6 @@ from torch import nn
 from torch._dynamo import config
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 config.cache_size_limit = 128

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_bf16_storage.py

@@ -5,7 +5,6 @@ from torch import nn
 from torch._dynamo import config
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 config.cache_size_limit = 128

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_caution.py

@@ -9,7 +9,6 @@ from torch import nn
 from torch._dynamo import config
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 config.cache_size_limit = 128

heavyball-2.1.3/test/test_chainable_cpu.py

@@ -0,0 +1,65 @@
+import os
+
+import torch
+
+import heavyball.chainable as C
+import heavyball.utils
+
+os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
+heavyball.utils.compile_mode = None
+
+
+def _identity_update(state, group, update, grad, param):
+    return update
+
+
+def test_chain_applies_update_on_cpu():
+    param = [torch.nn.Parameter(torch.zeros(2))]
+    grad = [torch.ones(2)]
+    group = {"lr": 0.1, "caution": False, "weight_decay": 0.0}
+
+    with torch.no_grad():
+        C.chain(lambda _: {}, group, grad, param, _identity_update)
+
+    assert torch.allclose(param[0].detach(), torch.full((2,), -0.1))
+
+
+def test_branch_merges_multiple_paths():
+    def double(_, __, update, ___, ____):
+        return [u * 2 for u in update]
+
+    def negate(_, __, update, ___, ____):
+        return [u * -1 for u in update]
+
+    def merge_fn(outputs):
+        return [sum(vals) / len(vals) for vals in zip(*outputs)]
+
+    branch = C.Branch([[double], [negate]], merge_fn)
+
+    update = [torch.ones(2)]
+    grad = [torch.ones(2)]
+    param = [torch.nn.Parameter(torch.ones(2))]
+
+    result = branch(lambda _: {}, {}, update, grad, param)
+    expected = torch.full_like(update[0], 0.5)
+    assert torch.allclose(result[0], expected)
+
+
+def test_set_indices_assigns_transform_ids():
+    def base(_, __, update, ___, ____, buffer):
+        assert buffer is not None
+        return update
+
+    zero_guard = C.ZeroGuard(base, ["buffer"])
+    assigned = C.set_indices([zero_guard], retain=False)[0]
+    assert assigned.transform_idx == 0
+
+    def state_fn(_x):
+        return {}
+
+    group = {"storage_dtype": "float32"}
+    update = [torch.ones(1)]
+    grad = [torch.ones(1)]
+    param = [torch.nn.Parameter(torch.ones(1))]
+
+    assigned(state_fn, group, update, grad, param)

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_channels_last.py

@@ -9,7 +9,6 @@ from torch import nn
 from torch._dynamo import config
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 heavyball.utils.zeroth_power_mode = "newtonschulz"

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_closure.py

@@ -6,7 +6,6 @@ from lightbench.utils import get_optim
 from torch import nn
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 

heavyball-2.1.3/test/test_cpu_features.py

@@ -0,0 +1,134 @@
+"""Fast CPU-only smoke tests for non-PSGD HeavyBall features."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+
+import pytest
+import torch
+from torch import nn
+
+import heavyball
+
+
+def _train_once(optimizer, model: nn.Module, data: torch.Tensor, target: torch.Tensor, steps: int = 3) -> float:
+    """Run a few optimization steps and return the final loss."""
+
+    loss = torch.tensor(float("nan"))
+    for _ in range(steps):
+        optimizer.zero_grad(set_to_none=True)
+        prediction = model(data)
+        loss = torch.nn.functional.mse_loss(prediction, target)
+        loss.backward()
+        optimizer.step()
+    return loss.item()
+
+
+def _parameter_drift(model: nn.Module, original: list[torch.Tensor]) -> float:
+    current = [param.detach() for param in model.parameters()]
+    diffs = [curr - init for curr, init in zip(current, original, strict=True)]
+    stacked = torch.cat([diff.reshape(-1) for diff in diffs])
+    return stacked.norm().item()
+
+
+def _make_batch(
+    in_features: int = 8, out_features: int = 4, batch: int = 16
+) -> tuple[nn.Module, torch.Tensor, torch.Tensor]:
+    torch.manual_seed(0x172893)
+    model = nn.Sequential(nn.Linear(in_features, out_features), nn.ReLU(), nn.Linear(out_features, out_features))
+    data = torch.randn(batch, in_features)
+    target = torch.randn(batch, out_features)
+    return model, data, target
+
+
+@pytest.mark.parametrize(
+    "opt_name",
+    [
+        "ForeachSOAP",
+        "Muon",
+        "ForeachAdamW",
+    ],
+)
+def test_selected_optimizers_run_on_cpu(opt_name: str) -> None:
+    model, data, target = _make_batch()
+    init = [param.detach().clone() for param in model.parameters()]
+
+    opt_cls = getattr(heavyball, opt_name)
+    optimizer = opt_cls(model.parameters(), warmup_steps=0)
+    final_loss = _train_once(optimizer, model, data, target, steps=3)
+
+    assert torch.isfinite(torch.tensor(final_loss))
+    assert _parameter_drift(model, init) > 0.0
+
+
+def test_caution_reduces_update_magnitude() -> None:
+    baseline_model, data, target = _make_batch()
+    cautious_model = deepcopy(baseline_model)
+
+    baseline_init = [param.detach().clone() for param in baseline_model.parameters()]
+    cautious_init = [param.detach().clone() for param in cautious_model.parameters()]
+
+    baseline_opt = heavyball.SGD(
+        baseline_model.parameters(),
+        lr=1e-3,
+        caution=False,
+    )
+    cautious_opt = heavyball.SGD(
+        cautious_model.parameters(),
+        lr=1e-3,
+        caution=True,
+    )
+
+    _train_once(baseline_opt, baseline_model, data, target)
+    _train_once(cautious_opt, cautious_model, data, target)
+
+    baseline_drift = _parameter_drift(baseline_model, baseline_init)
+    cautious_drift = _parameter_drift(cautious_model, cautious_init)
+
+    assert cautious_drift <= baseline_drift * 1.05  # caution should not overshoot compared to baseline
+
+
+def test_mars_flag_changes_behavior() -> None:
+    model_a, data, target = _make_batch()
+    model_b = deepcopy(model_a)
+
+    opt_a = heavyball.ForeachAdamW(model_a.parameters(), mars=False, warmup_steps=0)
+    opt_b = heavyball.ForeachAdamW(model_b.parameters(), mars=True, warmup_steps=0)
+
+    init = [param.detach().clone() for param in model_a.parameters()]
+
+    _train_once(opt_a, model_a, data, target)
+    _train_once(opt_b, model_b, data, target)
+
+    baseline_drift = _parameter_drift(model_a, init)
+    mars_drift = _parameter_drift(model_b, init)
+    assert baseline_drift > 0.0
+    assert mars_drift > 0.0
+
+    deltas = [a.detach() - b.detach() for a, b in zip(model_a.parameters(), model_b.parameters(), strict=True)]
+    combined = torch.cat([delta.reshape(-1) for delta in deltas])
+    assert combined.norm().item() > 1e-6  # mars path should diverge from baseline
+
+
+def test_sam_wrapper_requires_closure() -> None:
+    model = nn.Linear(4, 2)
+    base = heavyball.ForeachAdamW(model.parameters())
+    wrapper = heavyball.SAMWrapper(model.parameters(), wrapped_optimizer=base)
+
+    with pytest.raises(ValueError):
+        wrapper.step()
+
+    data = torch.randn(8, 4)
+    target = torch.randn(8, 2)
+
+    def closure():
+        wrapper.zero_grad()
+        loss = torch.nn.functional.mse_loss(model(data), target)
+        loss.backward()
+        return loss
+
+    before = [param.detach().clone() for param in model.parameters()]
+    wrapper.step(closure)
+    after = [param.detach() for param in model.parameters()]
+    diff = torch.cat([(a - b).reshape(-1) for a, b in zip(after, before, strict=True)])
+    assert diff.norm().item() > 0.0

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_ema.py

@@ -5,7 +5,6 @@ from torch import nn
 from torch._dynamo import config
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 config.cache_size_limit = 128

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_foreach.py

@@ -4,7 +4,6 @@ from lightbench.utils import get_optim
 from torch import nn
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 

heavyball-2.1.3/test/test_helpers_cpu.py

@@ -0,0 +1,107 @@
+import numpy as np
+import optuna
+import pandas as pd
+import torch
+from optuna.distributions import FloatDistribution, IntDistribution
+from optuna.samplers import RandomSampler
+from optuna.trial import TrialState
+
+from heavyball import helpers
+
+
+def test_bound_to_torch_roundtrip_cpu():
+    arr = np.arange(4, dtype=np.float64).reshape(2, 2)
+    tensor = helpers.bound_to_torch(arr.tobytes(), arr.shape, "cpu")
+    assert torch.allclose(tensor, torch.from_numpy(arr.T))
+
+
+def test_nextafter_matches_numpy():
+    forward = helpers.nextafter(0.5, 1.0)
+    backward = helpers.nextafter(1, 0)
+    assert forward == np.nextafter(0.5, 1.0)
+    assert backward == np.nextafter(1, 0)
+
+
+def test_untransform_numerical_param_torch_handles_steps():
+    dist = FloatDistribution(0.0, 1.0, step=0.1)
+    value = torch.tensor(0.46)
+    untransformed = helpers._untransform_numerical_param_torch(value, dist, transform_log=False)
+    assert torch.isclose(untransformed, torch.tensor(0.5))
+
+
+def test_simple_api_sampler_suggest_all_returns_expected():
+    distributions = {"x": FloatDistribution(0.0, 1.0), "y": IntDistribution(0, 3, step=1)}
+
+    class _Sampler(helpers.SimpleAPIBaseSampler):
+        def infer_relative_search_space(self, study, trial):
+            return self.search_space
+
+        def sample_relative(self, study, trial, search_space):
+            return {}
+
+        def sample_independent(self, study, trial, param_name, param_distribution):
+            return trial.params[param_name]
+
+    sampler = _Sampler(distributions)
+
+    class DummyTrial:
+        def __init__(self, params):
+            self.params = params
+
+        def _suggest(self, name, dist):
+            return self.params[name]
+
+    trial = DummyTrial({"x": 0.25, "y": 2})
+    suggestions = sampler.suggest_all(trial)
+    assert suggestions == {"x": 0.25, "y": 2}
+
+
+def test_botorch_sampler_sample_relative_smoke(monkeypatch):
+    search_space = {"width": FloatDistribution(0.0, 1.0)}
+    study = optuna.create_study(direction="minimize", sampler=RandomSampler(seed=0))
+    for _ in range(3):
+        trial = study.ask()
+        width = trial.suggest_float("width", 0.0, 1.0)
+        study.tell(trial, width)
+
+    sampler = helpers.BoTorchSampler(search_space, n_startup_trials=1, seed=0, device="cpu")
+
+    def _dummy_candidates(params, values, *_args):
+        assert params.shape[1] == 1
+        return params.mean(dim=0)
+
+    sampler._candidates_func = _dummy_candidates
+
+    pending = study.ask()
+    suggestion = sampler.sample_relative(study, pending, search_space)
+    assert "width" in suggestion
+    assert 0.0 <= suggestion["width"] <= 1.0
+
+
+def test_hebo_sampler_observe_and_sample(monkeypatch):
+    class DummyHEBO:
+        def __init__(self, *_args, **_kwargs):
+            self.observed = None
+
+        def suggest(self):
+            return pd.DataFrame([{"depth": 0.0}])
+
+        def observe(self, params, values):
+            self.observed = (params, values)
+
+    monkeypatch.setattr(helpers, "HEBO", DummyHEBO)
+
+    search_space = {"depth": FloatDistribution(0.0, 1.0)}
+    sampler = helpers.HEBOSampler(search_space, seed=1)
+
+    study = optuna.create_study(direction="minimize", sampler=RandomSampler(seed=1))
+    trial = study.ask()
+    trial.suggest_float("depth", 0.0, 1.0)
+    study.tell(trial, 0.2)
+
+    suggestion = sampler.sample_relative(study, study.ask(), search_space)
+    assert suggestion["depth"] == 0.0
+
+    completed = study.get_trials(deepcopy=False)[0]
+    sampler.after_trial(study, completed, TrialState.COMPLETE, [0.2])
+    assert sampler._hebo.observed is not None

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_hook.py

@@ -9,7 +9,6 @@ from torch import nn
 from torch._dynamo import config
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, hook_optimizer_into_model, set_torch
 
 heavyball.utils.compile_mode = "default"

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_mars.py

@@ -5,7 +5,6 @@ from torch import nn
 from torch._dynamo import config
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 config.cache_size_limit = 128

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_memory.py

@@ -4,7 +4,6 @@ from lightbench.utils import get_optim
 from torch import nn
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_memory_leak.py

@@ -6,7 +6,6 @@ from torch import nn
 from torch.nn import functional as F
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_merge.py

@@ -6,7 +6,6 @@ from lightbench.utils import get_optim
 from torch import nn
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import clean, set_torch
 
 

{heavyball-2.1.2 → heavyball-2.1.3}/test/test_nd_param.py

@@ -5,7 +5,6 @@ from torch import nn
 from torch._dynamo import config
 
 import heavyball
-import heavyball.utils
 from heavyball.utils import set_torch
 
 config.cache_size_limit = 2**20