"""Shared base class.

``torch.optim.Optimizer`` is used strictly as a base class, as the assignment
allows: it supplies parameter-group bookkeeping and the ``state`` dict, and
every update rule below is written out by hand.
"""

import torch


class BaseOptimizer(torch.optim.Optimizer):
    @staticmethod
    def _closure(closure):
        if closure is None:
            return None
        with torch.enable_grad():
            return closure()

    def memory_bytes(self) -> int:
        """Bytes held in optimizer state -- the number that matters for the
        memory-efficient category."""
        total = 0
        for st in self.state.values():
            for v in st.values():
                if torch.is_tensor(v):
                    total += v.numel() * v.element_size()
        return total