open-lars 0.1.0__tar.gz

open_lars-0.1.0/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 Ian Yang
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

open_lars-0.1.0/PKG-INFO

@@ -0,0 +1,132 @@
+Metadata-Version: 2.4
+Name: open-lars
+Version: 0.1.0
+Summary: A pure-PyTorch LARS optimizer, compatible with modern CUDA
+Author-email: Ian Yang <ian16yang@gmail.com>
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/GrignardReagent/open_lars
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=1.8
+Provides-Extra: test
+Requires-Dist: pytest>=4; extra == "test"
+Dynamic: license-file
+
+# open_lars
+
+A **pure-PyTorch implementation of the LARS optimizer** (Layer-wise Adaptive
+Rate Scaling; [You, Gitman & Ginsburg 2017](https://arxiv.org/abs/1708.03888)).
+
+LARS is the optimizer used by the SimCLR / SupCon lineage (Chen et al. 2020;
+Khosla et al. 2020) for large-batch contrastive pretraining: at very large
+batch sizes a single global learning rate is simultaneously too large for some
+layers and too small for others, so LARS rescales each layer's effective LR by
+a "trust ratio" - the ratio of that layer's weight norm to its gradient norm.
+
+## Why pure PyTorch?
+
+Older LARS packages computed the trust ratio in custom C++/CUDA kernels
+compiled at install time. That approach ages badly: it needs an `nvcc`
+matching the exact CUDA version your PyTorch wheel was built against (fragile
+on shared or HPC machines, impossible where no CUDA toolkit is installed), it
+pins you to one torch ABI so every torch upgrade means a rebuild, and once the
+underlying C++/ATen APIs churn, the extension stops building altogether.
+
+The trust ratio is a handful of scalar operations, so `open_lars` implements
+it as plain tensor ops instead. There is nothing to compile: it runs on CPU
+and on **any CUDA version supported by your PyTorch build** (also ROCm/MPS),
+with any `torch >= 1.8`.
+
+## Install
+
+```bash
+pip install open_lars
+```
+
+Or from source:
+
+```bash
+pip install .            # from a checkout
+pip install -e .         # editable, for development
+```
+
+The only dependency is `torch`. Once installed in an environment, it is
+importable from any directory you work in.
+
+## Usage
+
+`open_lars.LARS` is a standalone optimizer (not a wrapper around another
+optimizer), with SGD-with-momentum as the inner update rule, exactly as in
+the paper's Algorithm 1:
+
+```python
+from open_lars import LARS
+
+optimizer = LARS(model.parameters(), lr=1.6, momentum=0.9,
+                 weight_decay=1e-4, trust_coefficient=0.001)
+
+for x, y in loader:
+    optimizer.zero_grad()
+    loss = criterion(model(x), y)
+    loss.backward()
+    optimizer.step()
+```
+
+It is a standard `torch.optim.Optimizer`: parameter groups, `state_dict()` /
+`load_state_dict()` checkpointing, closures, and LR schedulers (e.g.
+linear-warmup + cosine decay, the schedule conventionally paired with LARS)
+all work as usual.
+
+### Update rule
+
+```text
+local_lr = trust_coefficient * ||w|| / (||∇L(w)|| + weight_decay * ||w|| + eps)
+v_t      = momentum * v_{t-1} + lr * local_lr * (∇L(w) + weight_decay * w)
+w_{t+1}  = w_t - v_t
+```
+
+Following standard practice (the paper; SimCLR/BYOL reference code), the trust
+ratio is **skipped for 1-D parameters** (biases, LayerNorm/BatchNorm weights),
+which fall back to plain SGD+momentum — applying per-layer adaptation to those
+is known to hurt convergence. If `||w||` or `||∇L(w)||` is zero, `local_lr`
+falls back to `1.0` so the parameter still receives a plain update.
+
+The trust-ratio computation is also exposed directly as a pure-torch function:
+
+```python
+from open_lars import compute_adaptive_lr
+```
+
+## Tests
+
+```bash
+pip install -e '.[test]'
+pytest
+```
+
+The suite covers the trust-ratio formula and its zero-norm guards across
+dtypes and devices (including half precision), the optimizer contract
+(pickling, checkpointing, schedulers, parameter groups, closures), and an
+equivalence suite verifying that `open_lars.LARS` reproduces the reference
+implementation in `tests/reference_lars.py` **bit-for-bit** over multi-step
+training trajectories, on CPU and CUDA, across momentum/weight-decay
+configurations. CUDA-specific tests skip automatically on CPU-only machines.
+
+## References
+
+- You, Gitman & Ginsburg, *Large Batch Training of Convolutional Networks*,
+  2017. <https://arxiv.org/abs/1708.03888>
+- Chen et al., *A Simple Framework for Contrastive Learning of Visual
+  Representations* (SimCLR), 2020. <https://arxiv.org/abs/2002.05709>
+- Khosla et al., *Supervised Contrastive Learning*, 2020.
+  <https://arxiv.org/abs/2004.11362>
+
+## License
+
+Apache License 2.0 — see `LICENSE`.

open_lars-0.1.0/README.md

@@ -0,0 +1,112 @@
+# open_lars
+
+A **pure-PyTorch implementation of the LARS optimizer** (Layer-wise Adaptive
+Rate Scaling; [You, Gitman & Ginsburg 2017](https://arxiv.org/abs/1708.03888)).
+
+LARS is the optimizer used by the SimCLR / SupCon lineage (Chen et al. 2020;
+Khosla et al. 2020) for large-batch contrastive pretraining: at very large
+batch sizes a single global learning rate is simultaneously too large for some
+layers and too small for others, so LARS rescales each layer's effective LR by
+a "trust ratio" - the ratio of that layer's weight norm to its gradient norm.
+
+## Why pure PyTorch?
+
+Older LARS packages computed the trust ratio in custom C++/CUDA kernels
+compiled at install time. That approach ages badly: it needs an `nvcc`
+matching the exact CUDA version your PyTorch wheel was built against (fragile
+on shared or HPC machines, impossible where no CUDA toolkit is installed), it
+pins you to one torch ABI so every torch upgrade means a rebuild, and once the
+underlying C++/ATen APIs churn, the extension stops building altogether.
+
+The trust ratio is a handful of scalar operations, so `open_lars` implements
+it as plain tensor ops instead. There is nothing to compile: it runs on CPU
+and on **any CUDA version supported by your PyTorch build** (also ROCm/MPS),
+with any `torch >= 1.8`.
+
+## Install
+
+```bash
+pip install open_lars
+```
+
+Or from source:
+
+```bash
+pip install .            # from a checkout
+pip install -e .         # editable, for development
+```
+
+The only dependency is `torch`. Once installed in an environment, it is
+importable from any directory you work in.
+
+## Usage
+
+`open_lars.LARS` is a standalone optimizer (not a wrapper around another
+optimizer), with SGD-with-momentum as the inner update rule, exactly as in
+the paper's Algorithm 1:
+
+```python
+from open_lars import LARS
+
+optimizer = LARS(model.parameters(), lr=1.6, momentum=0.9,
+                 weight_decay=1e-4, trust_coefficient=0.001)
+
+for x, y in loader:
+    optimizer.zero_grad()
+    loss = criterion(model(x), y)
+    loss.backward()
+    optimizer.step()
+```
+
+It is a standard `torch.optim.Optimizer`: parameter groups, `state_dict()` /
+`load_state_dict()` checkpointing, closures, and LR schedulers (e.g.
+linear-warmup + cosine decay, the schedule conventionally paired with LARS)
+all work as usual.
+
+### Update rule
+
+```text
+local_lr = trust_coefficient * ||w|| / (||∇L(w)|| + weight_decay * ||w|| + eps)
+v_t      = momentum * v_{t-1} + lr * local_lr * (∇L(w) + weight_decay * w)
+w_{t+1}  = w_t - v_t
+```
+
+Following standard practice (the paper; SimCLR/BYOL reference code), the trust
+ratio is **skipped for 1-D parameters** (biases, LayerNorm/BatchNorm weights),
+which fall back to plain SGD+momentum — applying per-layer adaptation to those
+is known to hurt convergence. If `||w||` or `||∇L(w)||` is zero, `local_lr`
+falls back to `1.0` so the parameter still receives a plain update.
+
+The trust-ratio computation is also exposed directly as a pure-torch function:
+
+```python
+from open_lars import compute_adaptive_lr
+```
+
+## Tests
+
+```bash
+pip install -e '.[test]'
+pytest
+```
+
+The suite covers the trust-ratio formula and its zero-norm guards across
+dtypes and devices (including half precision), the optimizer contract
+(pickling, checkpointing, schedulers, parameter groups, closures), and an
+equivalence suite verifying that `open_lars.LARS` reproduces the reference
+implementation in `tests/reference_lars.py` **bit-for-bit** over multi-step
+training trajectories, on CPU and CUDA, across momentum/weight-decay
+configurations. CUDA-specific tests skip automatically on CPU-only machines.
+
+## References
+
+- You, Gitman & Ginsburg, *Large Batch Training of Convolutional Networks*,
+  2017. <https://arxiv.org/abs/1708.03888>
+- Chen et al., *A Simple Framework for Contrastive Learning of Visual
+  Representations* (SimCLR), 2020. <https://arxiv.org/abs/2002.05709>
+- Khosla et al., *Supervised Contrastive Learning*, 2020.
+  <https://arxiv.org/abs/2004.11362>
+
+## License
+
+Apache License 2.0 — see `LICENSE`.

open_lars-0.1.0/open_lars/__init__.py

@@ -0,0 +1,5 @@
+from open_lars.__version__ import __version__  # noqa
+from open_lars.functional import compute_adaptive_lr
+from open_lars.lars import LARS
+
+__all__ = ["LARS", "compute_adaptive_lr"]

open_lars-0.1.0/open_lars/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

open_lars-0.1.0/open_lars/functional.py

@@ -0,0 +1,50 @@
+"""Functional core of the LARS optimizer.
+
+The trust ratio ("adaptive learning rate") is a handful of scalar operations,
+so it is expressed here as plain tensor ops: no compiled extension, so it runs
+on CPU or any CUDA/ROCm/MPS device torch itself supports, in any floating
+dtype, with no compilation step at install time.
+"""
+
+import torch
+
+__all__ = ["compute_adaptive_lr"]
+
+
+def compute_adaptive_lr(param_norm, grad_norm, weight_decay, eps, trust_coef,
+                        out=None):
+    """Compute the LARS local ("adaptive") learning rate for one parameter.
+
+    Implements the trust ratio of You, Gitman & Ginsburg (2017)::
+
+        adaptive_lr = trust_coef * ||w|| / (||g|| + weight_decay * ||w|| + eps)
+
+    with two guard cases:
+
+    1. ``param_norm == 0`` (e.g. a parameter initialised to zero): the ratio
+       would be zero and freeze the parameter forever, so fall back to 1.0
+       (plain SGD-style update).
+    2. ``grad_norm == 0`` (no useful gradient signal): the ratio is
+       meaningless, so fall back to 1.0 and let the base update rule decide.
+
+    Args:
+        param_norm (Tensor): scalar tensor, ||w|| of the parameter.
+        grad_norm (Tensor): scalar tensor, ||g|| of its (raw) gradient.
+        weight_decay (float): weight-decay coefficient folded into the
+            denominator (the paper's ``beta * ||w||`` term).
+        eps (float): small constant for numerical stability.
+        trust_coef (float): the paper's trust coefficient ``eta``.
+        out (Tensor, optional): tensor to write the result into.
+
+    Returns:
+        Tensor: scalar tensor holding the adaptive lr (``out`` if given).
+    """
+    adaptive_lr = torch.where(
+        (param_norm > 0) & (grad_norm > 0),
+        trust_coef * param_norm / (grad_norm + weight_decay * param_norm + eps),
+        torch.ones_like(param_norm),
+    )
+    if out is not None:
+        out.copy_(adaptive_lr)
+        return out
+    return adaptive_lr

open_lars-0.1.0/open_lars/lars.py

@@ -0,0 +1,128 @@
+"""
+LARS (Layer-wise Adaptive Rate Scaling) optimizer, in pure PyTorch.
+
+Reference
+---------
+You, Gitman & Ginsburg, "Large Batch Training of Convolutional Networks", 2017.
+https://arxiv.org/abs/1708.03888
+
+Why this exists: at very large batch sizes, a single global learning rate is
+simultaneously too large for some layers (divergence) and too small for
+others (slow convergence), because layers differ widely in parameter/gradient
+scale. LARS rescales each layer's *effective* LR by a "trust ratio" -- the
+ratio of that layer's weight norm to its gradient norm -- so every layer
+takes a step proportional to its own parameter scale rather than one shared
+LR. This is the optimizer Khosla et al. (SupCon) and the SimCLR lineage use
+for large-batch contrastive pretraining.
+
+Implements the paper's Algorithm 1 directly:
+    v_t   = momentum * v_{t-1}  +  lr * local_lr * (grad + weight_decay * w)
+    w_t+1 = w_t - v_t
+    local_lr = trust_coefficient * ||w|| / (||grad|| + weight_decay * ||w|| + eps)
+
+There is no custom CUDA kernel to compile at install time: the trust ratio is
+plain tensor ops (see ``open_lars.functional.compute_adaptive_lr``), so this
+runs against any current CUDA build of PyTorch -- or CPU/ROCm/MPS -- with no
+nvcc and no ABI pinning.
+
+Following standard practice (the paper; also SimCLR/BYOL reference code),
+LARS scaling is skipped for 1-D parameters (biases, LayerNorm/BatchNorm
+weights) -- applying a per-layer trust ratio to those is known to hurt
+convergence -- and falls back to plain SGD+momentum for them instead.
+"""
+
+import torch
+from torch.optim import Optimizer
+
+from open_lars.functional import compute_adaptive_lr
+
+__all__ = ["LARS"]
+
+
+class LARS(Optimizer):
+    """LARS as a self-contained :class:`torch.optim.Optimizer`.
+
+    LARS *is* the optimizer here (rather than a wrapper around another one),
+    with SGD-with-momentum as the inner update rule, exactly as in the
+    paper's Algorithm 1.
+
+    Args:
+        params: iterable of parameters or dicts defining parameter groups.
+        lr (float): global learning rate.
+        momentum (float, optional): SGD-style momentum used inside LARS.
+        weight_decay (float, optional): L2 penalty, folded into both the
+            trust-ratio denominator and the update direction.
+        trust_coefficient (float, optional): the paper's ``eta``.
+        eps (float, optional): numerical-stability constant in the
+            trust-ratio denominator.
+
+    Example::
+
+        optimizer = LARS(model.parameters(), lr=1.6, momentum=0.9,
+                         weight_decay=1e-4, trust_coefficient=0.001)
+
+        loss = loss_fn(model(input), target)
+        loss.backward()
+        optimizer.step()
+    """
+
+    def __init__(self, params, lr, momentum=0.9, weight_decay=0.0,
+                 trust_coefficient=0.001, eps=1e-8):
+        if lr < 0.0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if momentum < 0.0:
+            raise ValueError(f"Invalid momentum value: {momentum}")
+        if weight_decay < 0.0:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+        if trust_coefficient < 0.0:
+            raise ValueError(f"Invalid trust_coefficient value: {trust_coefficient}")
+        if eps < 0.0:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+        defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay,
+                        trust_coefficient=trust_coefficient, eps=eps)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            weight_decay = group["weight_decay"]
+            momentum = group["momentum"]
+            trust_coefficient = group["trust_coefficient"]
+            eps = group["eps"]
+            global_lr = group["lr"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError("LARS does not support sparse gradients")
+
+                # Trust ratio uses the RAW gradient norm in its denominator
+                # (matches the paper's ||∇L(w)|| term) -- weight decay is
+                # folded in separately, only for the velocity update itself.
+                if p.ndim > 1:
+                    local_lr = compute_adaptive_lr(
+                        torch.norm(p), torch.norm(grad),
+                        weight_decay, eps, trust_coefficient).item()
+                else:
+                    local_lr = 1.0  # biases / norm params: plain SGD+momentum, no trust ratio
+
+                d_p = grad if weight_decay == 0 else grad.add(p, alpha=weight_decay)
+                actual_lr = global_lr * local_lr
+
+                state = self.state[p]
+                if "momentum_buffer" not in state:
+                    buf = state["momentum_buffer"] = d_p.mul(actual_lr)
+                else:
+                    buf = state["momentum_buffer"]
+                    buf.mul_(momentum).add_(d_p, alpha=actual_lr)
+
+                p.add_(buf, alpha=-1.0)
+
+        return loss

open_lars-0.1.0/open_lars.egg-info/PKG-INFO

@@ -0,0 +1,132 @@
+Metadata-Version: 2.4
+Name: open-lars
+Version: 0.1.0
+Summary: A pure-PyTorch LARS optimizer, compatible with modern CUDA
+Author-email: Ian Yang <ian16yang@gmail.com>
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/GrignardReagent/open_lars
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=1.8
+Provides-Extra: test
+Requires-Dist: pytest>=4; extra == "test"
+Dynamic: license-file
+
+# open_lars
+
+A **pure-PyTorch implementation of the LARS optimizer** (Layer-wise Adaptive
+Rate Scaling; [You, Gitman & Ginsburg 2017](https://arxiv.org/abs/1708.03888)).
+
+LARS is the optimizer used by the SimCLR / SupCon lineage (Chen et al. 2020;
+Khosla et al. 2020) for large-batch contrastive pretraining: at very large
+batch sizes a single global learning rate is simultaneously too large for some
+layers and too small for others, so LARS rescales each layer's effective LR by
+a "trust ratio" - the ratio of that layer's weight norm to its gradient norm.
+
+## Why pure PyTorch?
+
+Older LARS packages computed the trust ratio in custom C++/CUDA kernels
+compiled at install time. That approach ages badly: it needs an `nvcc`
+matching the exact CUDA version your PyTorch wheel was built against (fragile
+on shared or HPC machines, impossible where no CUDA toolkit is installed), it
+pins you to one torch ABI so every torch upgrade means a rebuild, and once the
+underlying C++/ATen APIs churn, the extension stops building altogether.
+
+The trust ratio is a handful of scalar operations, so `open_lars` implements
+it as plain tensor ops instead. There is nothing to compile: it runs on CPU
+and on **any CUDA version supported by your PyTorch build** (also ROCm/MPS),
+with any `torch >= 1.8`.
+
+## Install
+
+```bash
+pip install open_lars
+```
+
+Or from source:
+
+```bash
+pip install .            # from a checkout
+pip install -e .         # editable, for development
+```
+
+The only dependency is `torch`. Once installed in an environment, it is
+importable from any directory you work in.
+
+## Usage
+
+`open_lars.LARS` is a standalone optimizer (not a wrapper around another
+optimizer), with SGD-with-momentum as the inner update rule, exactly as in
+the paper's Algorithm 1:
+
+```python
+from open_lars import LARS
+
+optimizer = LARS(model.parameters(), lr=1.6, momentum=0.9,
+                 weight_decay=1e-4, trust_coefficient=0.001)
+
+for x, y in loader:
+    optimizer.zero_grad()
+    loss = criterion(model(x), y)
+    loss.backward()
+    optimizer.step()
+```
+
+It is a standard `torch.optim.Optimizer`: parameter groups, `state_dict()` /
+`load_state_dict()` checkpointing, closures, and LR schedulers (e.g.
+linear-warmup + cosine decay, the schedule conventionally paired with LARS)
+all work as usual.
+
+### Update rule
+
+```text
+local_lr = trust_coefficient * ||w|| / (||∇L(w)|| + weight_decay * ||w|| + eps)
+v_t      = momentum * v_{t-1} + lr * local_lr * (∇L(w) + weight_decay * w)
+w_{t+1}  = w_t - v_t
+```
+
+Following standard practice (the paper; SimCLR/BYOL reference code), the trust
+ratio is **skipped for 1-D parameters** (biases, LayerNorm/BatchNorm weights),
+which fall back to plain SGD+momentum — applying per-layer adaptation to those
+is known to hurt convergence. If `||w||` or `||∇L(w)||` is zero, `local_lr`
+falls back to `1.0` so the parameter still receives a plain update.
+
+The trust-ratio computation is also exposed directly as a pure-torch function:
+
+```python
+from open_lars import compute_adaptive_lr
+```
+
+## Tests
+
+```bash
+pip install -e '.[test]'
+pytest
+```
+
+The suite covers the trust-ratio formula and its zero-norm guards across
+dtypes and devices (including half precision), the optimizer contract
+(pickling, checkpointing, schedulers, parameter groups, closures), and an
+equivalence suite verifying that `open_lars.LARS` reproduces the reference
+implementation in `tests/reference_lars.py` **bit-for-bit** over multi-step
+training trajectories, on CPU and CUDA, across momentum/weight-decay
+configurations. CUDA-specific tests skip automatically on CPU-only machines.
+
+## References
+
+- You, Gitman & Ginsburg, *Large Batch Training of Convolutional Networks*,
+  2017. <https://arxiv.org/abs/1708.03888>
+- Chen et al., *A Simple Framework for Contrastive Learning of Visual
+  Representations* (SimCLR), 2020. <https://arxiv.org/abs/2002.05709>
+- Khosla et al., *Supervised Contrastive Learning*, 2020.
+  <https://arxiv.org/abs/2004.11362>
+
+## License
+
+Apache License 2.0 — see `LICENSE`.

open_lars-0.1.0/open_lars.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+LICENSE
+README.md
+pyproject.toml
+setup.py
+open_lars/__init__.py
+open_lars/__version__.py
+open_lars/functional.py
+open_lars/lars.py
+open_lars.egg-info/PKG-INFO
+open_lars.egg-info/SOURCES.txt
+open_lars.egg-info/dependency_links.txt
+open_lars.egg-info/requires.txt
+open_lars.egg-info/top_level.txt
+tests/test_compute_adaptive_lr.py
+tests/test_lars.py
+tests/test_pickle.py

open_lars-0.1.0/open_lars.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

open_lars-0.1.0/open_lars.egg-info/requires.txt

@@ -0,0 +1,4 @@
+torch>=1.8
+
+[test]
+pytest>=4

open_lars-0.1.0/open_lars.egg-info/top_level.txt

@@ -0,0 +1 @@
+open_lars

open_lars-0.1.0/pyproject.toml

@@ -0,0 +1,35 @@
+[build-system]
+requires = ["setuptools>=64"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "open-lars"
+description = "A pure-PyTorch LARS optimizer, compatible with modern CUDA"
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "Apache-2.0" }
+authors = [{ name = "Ian Yang", email = "ian16yang@gmail.com" }]
+dependencies = ["torch>=1.8"]
+dynamic = ["version"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3 :: Only",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+[project.urls]
+Homepage = "https://github.com/GrignardReagent/open_lars"
+
+[project.optional-dependencies]
+test = ["pytest>=4"]
+
+[tool.setuptools]
+packages = ["open_lars"]
+
+[tool.setuptools.dynamic]
+version = { attr = "open_lars.__version__.__version__" }
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

open_lars-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

open_lars-0.1.0/setup.py

@@ -0,0 +1,9 @@
+"""Shim for pip versions whose PEP 660 editable-install support is broken
+(e.g. Ubuntu 22.04's stock pip 22.0.2): with this file present they fall
+back to the legacy `setup.py develop` path for `pip install -e .`.
+All metadata lives in pyproject.toml; modern pip ignores this file.
+"""
+
+from setuptools import setup
+
+setup()

open_lars-0.1.0/tests/test_compute_adaptive_lr.py

@@ -0,0 +1,104 @@
+"""Tests for open_lars.compute_adaptive_lr.
+
+The function is pure tensor ops, so every case is checked across floating
+dtypes and (when available) across devices, including half precision.
+"""
+
+import pytest
+import torch
+
+from open_lars import compute_adaptive_lr
+
+DEVICES = ['cpu'] + (['cuda'] if torch.cuda.is_available() else [])
+DTYPES = [torch.float, torch.double]
+
+
+def expected_ratio(param_norm, grad_norm, weight_decay, eps, trust_coef):
+    return trust_coef * param_norm / (grad_norm + weight_decay * param_norm + eps)
+
+
+@pytest.mark.parametrize('device', DEVICES)
+@pytest.mark.parametrize('dtype', DTYPES)
+@pytest.mark.parametrize('param_norm,grad_norm,weight_decay,eps,trust_coef', [
+    (2.5, 0.5, 1e-4, 1e-8, 0.001),
+    (0.037, 12.0, 1e-4, 1e-8, 0.001),
+    (1.0, 1.0, 1.0, 2.0, 1.0),
+    (3.0, 0.25, 0.0, 0.0, 0.02),
+])
+def test_matches_formula(device, dtype, param_norm, grad_norm, weight_decay,
+                         eps, trust_coef):
+    result = compute_adaptive_lr(
+        torch.tensor(param_norm, dtype=dtype, device=device),
+        torch.tensor(grad_norm, dtype=dtype, device=device),
+        weight_decay, eps, trust_coef)
+
+    expected = expected_ratio(param_norm, grad_norm, weight_decay, eps, trust_coef)
+    assert torch.allclose(result, torch.tensor(expected, dtype=dtype, device=device))
+
+
+@pytest.mark.parametrize('device', DEVICES)
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_zero_param_norm_falls_back_to_one(device, dtype):
+    # ||w|| == 0 would give a zero ratio and freeze the parameter forever.
+    result = compute_adaptive_lr(
+        torch.tensor(0., dtype=dtype, device=device),
+        torch.tensor(1., dtype=dtype, device=device),
+        1e-4, 1e-8, 0.001)
+
+    assert result == torch.tensor(1., dtype=dtype, device=device)
+
+
+@pytest.mark.parametrize('device', DEVICES)
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_zero_grad_norm_falls_back_to_one(device, dtype):
+    # ||g|| == 0 carries no signal; defer to the base update rule.
+    result = compute_adaptive_lr(
+        torch.tensor(1., dtype=dtype, device=device),
+        torch.tensor(0., dtype=dtype, device=device),
+        1e-4, 1e-8, 0.001)
+
+    assert result == torch.tensor(1., dtype=dtype, device=device)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='cuda required')
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_cpu_and_gpu_agree(dtype):
+    args = (0.8, 3.2, 1e-4, 1e-8, 0.001)
+
+    cpu = compute_adaptive_lr(
+        torch.tensor(args[0], dtype=dtype),
+        torch.tensor(args[1], dtype=dtype), *args[2:])
+    gpu = compute_adaptive_lr(
+        torch.tensor(args[0], dtype=dtype, device='cuda'),
+        torch.tensor(args[1], dtype=dtype, device='cuda'), *args[2:])
+
+    assert torch.allclose(cpu, gpu.cpu())
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='cuda required')
+@pytest.mark.parametrize('param_norm,grad_norm,expected', [
+    (2.5, 0.5, 0.001 * 2.5 / (0.5 + 1e-4 * 2.5 + 1e-8)),
+    (0.0, 1.0, 1.0),
+    (1.0, 0.0, 1.0),
+])
+def test_half_precision_on_gpu(param_norm, grad_norm, expected):
+    result = compute_adaptive_lr(
+        torch.tensor(param_norm, dtype=torch.half, device='cuda'),
+        torch.tensor(grad_norm, dtype=torch.half, device='cuda'),
+        1e-4, 1e-8, 0.001)
+
+    assert torch.allclose(result, torch.tensor(expected, dtype=torch.half, device='cuda'))
+
+
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_out_argument_is_written_in_place(dtype):
+    out = torch.tensor(0., dtype=dtype)
+
+    returned = compute_adaptive_lr(
+        torch.tensor(2.5, dtype=dtype),
+        torch.tensor(0.5, dtype=dtype),
+        1e-4, 1e-8, 0.001, out)
+
+    assert returned is out
+    expected = expected_ratio(2.5, 0.5, 1e-4, 1e-8, 0.001)
+    assert torch.allclose(out, torch.tensor(expected, dtype=dtype))

open_lars-0.1.0/tests/test_lars.py

@@ -0,0 +1,284 @@
+"""Behavioural tests for open_lars.LARS.
+
+The centrepiece is the equivalence suite: open_lars.LARS must produce
+*bit-identical* parameter trajectories to the reference implementation used
+by IY036 (vendored verbatim in tests/reference_lars.py), on CPU and CUDA,
+with and without weight decay / momentum, for both >1-D parameters (trust
+ratio applied) and 1-D parameters (plain SGD+momentum fallback).
+"""
+
+import copy
+import math
+
+import pytest
+import torch
+from torch import nn
+
+from open_lars import LARS
+from tests.reference_lars import LARS as ReferenceLARS
+
+DEVICES = ['cpu'] + (['cuda'] if torch.cuda.is_available() else [])
+
+
+def make_model(device):
+    # Mix of 2-D weights (trust ratio applied), biases and LayerNorm weights
+    # (1-D: plain SGD+momentum), mirroring the transformer trained in IY036.
+    return nn.Sequential(
+        nn.Linear(8, 16),
+        nn.LayerNorm(16),
+        nn.ReLU(),
+        nn.Linear(16, 4),
+    ).to(device)
+
+
+def train_steps(model, optimizer, device, n_steps=10, scheduler=None, seed=0):
+    torch.manual_seed(seed)
+    for _ in range(n_steps):
+        x = torch.randn(32, 8, device=device)
+        y = torch.randn(32, 4, device=device)
+        optimizer.zero_grad()
+        loss = nn.functional.mse_loss(model(x), y)
+        loss.backward()
+        optimizer.step()
+        if scheduler is not None:
+            scheduler.step()
+    return loss
+
+
+# ── Equivalence with the IY036 reference implementation ──────────────────────
+
+@pytest.mark.parametrize('device', DEVICES)
+@pytest.mark.parametrize('momentum,weight_decay', [
+    (0.9, 1e-4),   # IY036's configuration
+    (0.9, 0.0),
+    (0.0, 1e-4),
+    (0.0, 0.0),
+])
+def test_matches_reference_implementation(device, momentum, weight_decay):
+    torch.manual_seed(42)
+    model_a = make_model(device)
+    model_b = copy.deepcopy(model_a)
+
+    opt_a = LARS(model_a.parameters(), lr=0.5, momentum=momentum,
+                 weight_decay=weight_decay, trust_coefficient=0.001)
+    opt_b = ReferenceLARS(model_b.parameters(), lr=0.5, momentum=momentum,
+                          weight_decay=weight_decay, trust_coefficient=0.001)
+
+    train_steps(model_a, opt_a, device, n_steps=10)
+    train_steps(model_b, opt_b, device, n_steps=10)
+
+    for p_a, p_b in zip(model_a.parameters(), model_b.parameters()):
+        assert torch.equal(p_a, p_b)
+
+
+# ── Update-rule unit checks ──────────────────────────────────────────────────
+
+def test_single_step_matches_formula_2d():
+    torch.manual_seed(0)
+    p = torch.randn(4, 4, requires_grad=True)
+    p_init = p.detach().clone()
+    lr, wd, tc, eps = 0.1, 1e-4, 0.001, 1e-8
+
+    loss = (p ** 2).sum()
+    loss.backward()
+    grad = p.grad.detach().clone()
+
+    LARS([p], lr=lr, momentum=0.9, weight_decay=wd, trust_coefficient=tc).step()
+
+    w_norm = p_init.norm()
+    g_norm = grad.norm()
+    local_lr = (tc * w_norm / (g_norm + wd * w_norm + eps)).item()
+    expected = p_init - lr * local_lr * (grad + wd * p_init)
+    assert torch.allclose(p.detach(), expected, atol=1e-7)
+
+
+def test_1d_params_use_plain_sgd_momentum():
+    # A bias-like 1-D parameter must be updated with local_lr == 1.0.
+    torch.manual_seed(0)
+    p = torch.randn(6, requires_grad=True)
+    p_init = p.detach().clone()
+    lr = 0.01
+
+    (p ** 2).sum().backward()
+    grad = p.grad.detach().clone()
+
+    LARS([p], lr=lr, momentum=0.9, weight_decay=0.0).step()
+
+    assert torch.allclose(p.detach(), p_init - lr * grad, atol=1e-7)
+
+
+def test_momentum_accumulates_velocity():
+    # Two steps with a constant gradient: v2 = m*v1 + lr*g, w2 = w0 - v1 - v2.
+    p = torch.ones(3, 3, requires_grad=True)
+    lr, m, tc, eps = 0.1, 0.9, 0.001, 1e-8
+    grad = torch.full((3, 3), 0.5)
+
+    opt = LARS([p], lr=lr, momentum=m, weight_decay=0.0, trust_coefficient=tc)
+
+    w = p.detach().clone()
+    v = torch.zeros_like(w)
+    for _ in range(2):
+        local_lr = (tc * w.norm() / (grad.norm() + eps)).item()
+        v = m * v + lr * local_lr * grad
+        w = w - v
+        p.grad = grad.clone()
+        opt.step()
+
+    assert torch.allclose(p.detach(), w, atol=1e-7)
+
+
+def test_zero_weight_param_still_updates():
+    # ||w|| == 0 would give trust ratio 0 and freeze the parameter forever;
+    # the guard must fall back to local_lr = 1.0 instead.
+    p = torch.zeros(4, 4, requires_grad=True)
+    p.grad = torch.ones(4, 4)
+
+    LARS([p], lr=0.1, momentum=0.0).step()
+
+    assert torch.allclose(p.detach(), torch.full((4, 4), -0.1))
+
+
+def test_zero_grad_is_a_noop():
+    p = torch.randn(4, 4, requires_grad=True)
+    p_init = p.detach().clone()
+    p.grad = torch.zeros(4, 4)
+
+    LARS([p], lr=0.1, momentum=0.9, weight_decay=0.0).step()
+
+    assert torch.equal(p.detach(), p_init)
+
+
+# ── torch.optim.Optimizer contract ───────────────────────────────────────────
+
+def test_closure_returns_loss():
+    p = torch.randn(3, 3, requires_grad=True)
+    opt = LARS([p], lr=0.1)
+
+    def closure():
+        opt.zero_grad()
+        loss = (p ** 2).sum()
+        loss.backward()
+        return loss
+
+    loss = opt.step(closure)
+    assert loss is not None and loss.item() >= 0.0
+
+
+@pytest.mark.parametrize('device', DEVICES)
+def test_state_dict_roundtrip(device):
+    torch.manual_seed(1)
+    model_a = make_model(device)
+    opt_a = LARS(model_a.parameters(), lr=0.5, weight_decay=1e-4)
+    train_steps(model_a, opt_a, device, n_steps=5)
+
+    # Resume from a checkpoint into fresh objects. The deepcopy stands in for
+    # a torch.save/torch.load round trip -- loading a live state_dict directly
+    # would alias opt_a's momentum buffers.
+    model_b = copy.deepcopy(model_a)
+    opt_b = LARS(model_b.parameters(), lr=0.5, weight_decay=1e-4)
+    opt_b.load_state_dict(copy.deepcopy(opt_a.state_dict()))
+
+    # ...and both must continue on the exact same trajectory.
+    train_steps(model_a, opt_a, device, n_steps=5, seed=2)
+    train_steps(model_b, opt_b, device, n_steps=5, seed=2)
+
+    for p_a, p_b in zip(model_a.parameters(), model_b.parameters()):
+        assert torch.equal(p_a, p_b)
+
+
+def test_lr_scheduler_compatibility():
+    # IY036 pairs LARS with linear-warmup + cosine decay
+    # (transformers.get_cosine_schedule_with_warmup); reproduce that schedule
+    # with a plain LambdaLR to keep this package dependency-free.
+    torch.manual_seed(3)
+    model = make_model('cpu')
+    opt = LARS(model.parameters(), lr=1.0, weight_decay=1e-4)
+
+    total, warmup = 20, 5
+
+    def warmup_cosine(step):
+        if step < warmup:
+            return step / warmup
+        progress = (step - warmup) / (total - warmup)
+        return 0.5 * (1.0 + math.cos(math.pi * progress))
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(opt, warmup_cosine)
+
+    train_steps(model, opt, 'cpu', n_steps=total, scheduler=scheduler)
+    assert opt.param_groups[0]['lr'] == pytest.approx(0.0, abs=1e-6)
+
+
+def test_param_groups_with_distinct_hyperparams():
+    w = torch.randn(4, 4, requires_grad=True)
+    b = torch.randn(4, requires_grad=True)
+    opt = LARS([
+        {'params': [w], 'weight_decay': 1e-4},
+        {'params': [b], 'weight_decay': 0.0, 'lr': 0.01},
+    ], lr=0.5)
+
+    assert opt.param_groups[0]['weight_decay'] == 1e-4
+    assert opt.param_groups[1]['lr'] == 0.01
+
+    w.grad = torch.ones_like(w)
+    b.grad = torch.ones_like(b)
+    opt.step()  # must not raise
+
+
+def test_invalid_hyperparams_raise():
+    p = [torch.zeros(2, 2)]
+    with pytest.raises(ValueError):
+        LARS(p, lr=-0.1)
+    with pytest.raises(ValueError):
+        LARS(p, lr=0.1, momentum=-0.5)
+    with pytest.raises(ValueError):
+        LARS(p, lr=0.1, weight_decay=-1e-4)
+    with pytest.raises(ValueError):
+        LARS(p, lr=0.1, trust_coefficient=-0.001)
+    with pytest.raises(ValueError):
+        LARS(p, lr=0.1, eps=-1e-8)
+
+
+# ── End-to-end sanity ────────────────────────────────────────────────────────
+
+@pytest.mark.parametrize('device', DEVICES)
+def test_converges_on_toy_regression(device):
+    torch.manual_seed(7)
+    true_w = torch.randn(8, 1, device=device)
+    x = torch.randn(256, 8, device=device)
+    y = x @ true_w
+
+    model = nn.Linear(8, 1, bias=False).to(device)
+    opt = LARS(model.parameters(), lr=5.0, momentum=0.9,
+               trust_coefficient=0.01)
+
+    initial_loss = nn.functional.mse_loss(model(x), y).item()
+    for _ in range(200):
+        opt.zero_grad()
+        loss = nn.functional.mse_loss(model(x), y)
+        loss.backward()
+        opt.step()
+
+    assert loss.item() < 0.05 * initial_loss
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='cuda required')
+def test_cuda_matches_cpu_trajectory():
+    torch.manual_seed(11)
+    model_cpu = make_model('cpu')
+    model_gpu = copy.deepcopy(model_cpu).cuda()
+
+    opt_cpu = LARS(model_cpu.parameters(), lr=0.5, weight_decay=1e-4)
+    opt_gpu = LARS(model_gpu.parameters(), lr=0.5, weight_decay=1e-4)
+
+    torch.manual_seed(0)
+    x = torch.randn(32, 8)
+    y = torch.randn(32, 4)
+    for _ in range(5):
+        for model, opt, xd, yd in [(model_cpu, opt_cpu, x, y),
+                                   (model_gpu, opt_gpu, x.cuda(), y.cuda())]:
+            opt.zero_grad()
+            nn.functional.mse_loss(model(xd), yd).backward()
+            opt.step()
+
+    for p_c, p_g in zip(model_cpu.parameters(), model_gpu.parameters()):
+        assert torch.allclose(p_c, p_g.cpu(), atol=1e-5)

open_lars-0.1.0/tests/test_pickle.py

@@ -0,0 +1,36 @@
+"""LARS must survive (de)serialisation with its hyperparameters intact."""
+
+import copy
+import pickle
+
+import pytest
+import torch
+
+from open_lars import LARS
+
+
+@pytest.fixture
+def lars():
+    return LARS([torch.zeros(2, 2)], lr=0.1, momentum=0.8,
+                weight_decay=1e-4, trust_coefficient=0.42)
+
+
+def assert_hyperparams_preserved(opt):
+    assert isinstance(opt, LARS)
+    group = opt.param_groups[0]
+    assert group['lr'] == 0.1
+    assert group['momentum'] == 0.8
+    assert group['weight_decay'] == 1e-4
+    assert group['trust_coefficient'] == 0.42
+
+
+def test_pickle_roundtrip(lars):
+    assert_hyperparams_preserved(pickle.loads(pickle.dumps(lars)))
+
+
+def test_copy(lars):
+    assert_hyperparams_preserved(copy.copy(lars))
+
+
+def test_deepcopy(lars):
+    assert_hyperparams_preserved(copy.deepcopy(lars))