torchzero 0.3.9__py3-none-any.whl → 0.3.10__py3-none-any.whl

torchzero/modules/ops/reduce.py

@@ -5,7 +5,7 @@ from typing import Any, cast
 
 import torch
 
-from ...core import Chainable, Module, Target, Vars, maybe_chain
+from ...core import Chainable, Module, Target, Var, maybe_chain
 
 
 class ReduceOperation(Module, ABC):
@@ -26,25 +26,25 @@ class ReduceOperation(Module, ABC):
             raise ValueError('At least one operand must be a module')
 
     @abstractmethod
-    def transform(self, vars: Vars, *operands: Any | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *operands: Any | list[torch.Tensor]) -> list[torch.Tensor]:
         """applies the operation to operands"""
         raise NotImplementedError
 
     @torch.no_grad
-    def step(self, vars: Vars) -> Vars:
+    def step(self, var: Var) -> Var:
         # pass cloned update to all module operands
         processed_operands: list[Any | list[torch.Tensor]] = self.operands.copy()
 
         for i, v in enumerate(self.operands):
             if f'operand_{i}' in self.children:
                 v: Module
-                updated_vars = v.step(vars.clone(clone_update=True))
-                processed_operands[i] = updated_vars.get_update()
-                vars.update_attrs_from_clone_(updated_vars) # update loss, grad, etc if this module calculated them
+                updated_var = v.step(var.clone(clone_update=True))
+                processed_operands[i] = updated_var.get_update()
+                var.update_attrs_from_clone_(updated_var) # update loss, grad, etc if this module calculated them
 
-        transformed = self.transform(vars, *processed_operands)
-        vars.update = transformed
-        return vars
+        transformed = self.transform(var, *processed_operands)
+        var.update = transformed
+        return var
 
 class Sum(ReduceOperation):
     USE_MEAN = False
@@ -52,7 +52,7 @@ class Sum(ReduceOperation):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         sum = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:
@@ -76,9 +76,9 @@ class WeightedSum(ReduceOperation):
         super().__init__(defaults=defaults, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
-        weights = self.settings[vars.params[0]]['weights']
+        weights = self.settings[var.params[0]]['weights']
         sum = cast(list, sorted_inputs[0])
         torch._foreach_mul_(sum, weights[0])
         if len(sorted_inputs) > 1:
@@ -98,7 +98,7 @@ class Median(ReduceOperation):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         res = []
         lists = [i for i in inputs if isinstance(i, list)]
         floats = [i for i in inputs if isinstance(i, (int,float))]
@@ -111,7 +111,7 @@ class Prod(ReduceOperation):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         prod = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:
@@ -125,7 +125,7 @@ class MaximumModules(ReduceOperation):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         maximum = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:
@@ -139,7 +139,7 @@ class MinimumModules(ReduceOperation):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         minimum = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:

torchzero/modules/ops/split.py

@@ -3,46 +3,46 @@ from typing import cast
 
 import torch
 
-from ...core import Chainable, Module, Vars
+from ...core import Chainable, Module, Var
 
 
 def _split(
     module: Module,
     idxs,
     params,
-    vars: Vars,
+    var: Var,
 ):
     split_params = [p for i,p in enumerate(params) if i in idxs]
 
     split_grad = None
-    if vars.grad is not None:
-        split_grad = [g for i,g in enumerate(vars.grad) if i in idxs]
+    if var.grad is not None:
+        split_grad = [g for i,g in enumerate(var.grad) if i in idxs]
 
     split_update = None
-    if vars.update is not None:
-        split_update = [u for i,u in enumerate(vars.update) if i in idxs]
+    if var.update is not None:
+        split_update = [u for i,u in enumerate(var.update) if i in idxs]
 
-    split_vars = vars.clone(clone_update=False)
-    split_vars.params = split_params
-    split_vars.grad = split_grad
-    split_vars.update = split_update
+    split_var = var.clone(clone_update=False)
+    split_var.params = split_params
+    split_var.grad = split_grad
+    split_var.update = split_update
 
-    split_vars = module.step(split_vars)
+    split_var = module.step(split_var)
 
-    if (vars.grad is None) and (split_vars.grad is not None):
-        vars.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+    if (var.grad is None) and (split_var.grad is not None):
+        var.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
 
-    if split_vars.update is not None:
+    if split_var.update is not None:
 
-        if vars.update is None:
-            if vars.grad is None: vars.update = [cast(torch.Tensor, None) for _ in vars.params]
-            else: vars.update = [g.clone() for g in vars.grad]
+        if var.update is None:
+            if var.grad is None: var.update = [cast(torch.Tensor, None) for _ in var.params]
+            else: var.update = [g.clone() for g in var.grad]
 
-        for idx, u in zip(idxs, split_vars.update):
-            vars.update[idx] = u
+        for idx, u in zip(idxs, split_var.update):
+            var.update[idx] = u
 
-    vars.update_attrs_from_clone_(split_vars)
-    return vars
+    var.update_attrs_from_clone_(split_var)
+    return var
 
 class Split(Module):
     """Apply `true` modules to all parameters filtered by `filter`, apply `false` modules to all other parameters."""
@@ -53,9 +53,9 @@ class Split(Module):
         if true is not None: self.set_child('true', true)
         if false is not None: self.set_child('false', false)
 
-    def step(self, vars):
+    def step(self, var):
 
-        params = vars.params
+        params = var.params
         filter = self.settings[params[0]]['filter']
 
         true_idxs = []
@@ -66,10 +66,10 @@ class Split(Module):
 
         if 'true' in self.children:
             true = self.children['true']
-            vars = _split(true, idxs=true_idxs, params=params, vars=vars)
+            var = _split(true, idxs=true_idxs, params=params, var=var)
 
         if 'false' in self.children:
             false = self.children['false']
-            vars = _split(false, idxs=false_idxs, params=params, vars=vars)
+            var = _split(false, idxs=false_idxs, params=params, var=var)
 
-        return vars
+        return var

torchzero/modules/ops/switch.py

@@ -23,16 +23,16 @@ class Alternate(Module):
         self.global_state['steps_to_next'] = steps[0] if isinstance(steps, list) else steps
 
     @torch.no_grad
-    def step(self, vars):
+    def step(self, var):
         # get current module
         current_module_idx = self.global_state.setdefault('current_module_idx', 0)
         module = self.children[f'module_{current_module_idx}']
 
         # step
-        vars = module.step(vars.clone(clone_update=False))
+        var = module.step(var.clone(clone_update=False))
 
         # number of steps until next module
-        steps = self.settings[vars.params[0]]['steps']
+        steps = self.settings[var.params[0]]['steps']
         if isinstance(steps, int): steps = [steps]*len(self.children)
 
         if 'steps_to_next' not in self.global_state:
@@ -51,7 +51,7 @@ class Alternate(Module):
 
             self.global_state['steps_to_next'] = steps[self.global_state['current_module_idx']]
 
-        return vars
+        return var
 
 class Switch(Alternate):
     """switch to next module after some steps"""

torchzero/modules/ops/unary.py

@@ -3,7 +3,7 @@ from collections import deque
 import torch
 
 from ...core import TensorwiseTransform, Target, Transform
-from ...utils import TensorList
+from ...utils import TensorList, unpack_dicts,unpack_states
 
 class UnaryLambda(Transform):
     def __init__(self, fn, target: "Target" = 'update'):
@@ -11,8 +11,8 @@ class UnaryLambda(Transform):
         super().__init__(defaults=defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        return self.settings[params[0]]['fn'](tensors)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        return settings[0]['fn'](tensors)
 
 class UnaryParameterwiseLambda(TensorwiseTransform):
     def __init__(self, fn, target: "Target" = 'update'):
@@ -20,8 +20,8 @@ class UnaryParameterwiseLambda(TensorwiseTransform):
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def transform(self, tensor, param, grad, vars):
-        return self.settings[param]['fn'](tensor)
+    def apply_tensor(self, tensor, param, grad, loss, state, settings):
+        return settings['fn'](tensor)
 
 class CustomUnaryOperation(Transform):
     def __init__(self, name: str, target: "Target" = 'update'):
@@ -29,35 +29,35 @@ class CustomUnaryOperation(Transform):
         super().__init__(defaults=defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        return getattr(tensors, self.settings[params[0]]['name'])()
+    def apply(self, tensors, params, grads, loss, states, settings):
+        return getattr(tensors, settings[0]['name'])()
 
 
 class Abs(Transform):
     def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         torch._foreach_abs_(tensors)
         return tensors
 
 class Sign(Transform):
     def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         torch._foreach_sign_(tensors)
         return tensors
 
 class Exp(Transform):
     def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         torch._foreach_exp_(tensors)
         return tensors
 
 class Sqrt(Transform):
     def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         torch._foreach_sqrt_(tensors)
         return tensors
 
@@ -66,8 +66,8 @@ class Reciprocal(Transform):
         defaults = dict(eps = eps)
         super().__init__(defaults, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        eps = self.get_settings('eps', params=params)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        eps = [s['eps'] for s in settings]
         if any(e != 0 for e in eps): torch._foreach_add_(tensors, eps)
         torch._foreach_reciprocal_(tensors)
         return tensors
@@ -75,7 +75,7 @@ class Reciprocal(Transform):
 class Negate(Transform):
     def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         torch._foreach_neg_(tensors)
         return tensors
 
@@ -97,8 +97,8 @@ class NanToNum(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        nan, posinf, neginf = self.get_settings('nan', 'posinf', 'neginf', params=params)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        nan, posinf, neginf = unpack_dicts(settings, 'nan', 'posinf', 'neginf')
         return [t.nan_to_num_(nan_i, posinf_i, neginf_i) for t, nan_i, posinf_i, neginf_i in zip(tensors, nan, posinf, neginf)]
 
 class Rescale(Transform):
@@ -108,8 +108,8 @@ class Rescale(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        min,max = self.get_settings('min','max', params=params)
-        tensorwise = self.settings[params[0]]['tensorwise']
+    def apply(self, tensors, params, grads, loss, states, settings):
+        min, max = unpack_dicts(settings, 'min','max')
+        tensorwise = settings[0]['tensorwise']
         dim = None if tensorwise else 'global'
-        return TensorList(tensors).rescale(min=min, max=max, eps=self.settings[params[0]]['eps'], dim=dim)
+        return TensorList(tensors).rescale(min=min, max=max, eps=settings[0]['eps'], dim=dim)

torchzero/modules/ops/utility.py

@@ -9,47 +9,47 @@ from ...utils.tensorlist import Distributions, TensorList
 class Clone(Transform):
     def __init__(self): super().__init__({}, uses_grad=False)
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars): return [t.clone() for t in tensors]
+    def apply(self, tensors, params, grads, loss, states, settings): return [t.clone() for t in tensors]
 
 class Grad(Module):
     def __init__(self):
         super().__init__({})
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [g.clone() for g in vars.get_grad()]
-        return vars
+    def step(self, var):
+        var.update = [g.clone() for g in var.get_grad()]
+        return var
 
 class Params(Module):
     def __init__(self):
         super().__init__({})
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [p.clone() for p in vars.params]
-        return vars
+    def step(self, var):
+        var.update = [p.clone() for p in var.params]
+        return var
 
 class Update(Module):
     def __init__(self):
         super().__init__({})
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [u.clone() for u in vars.get_update()]
-        return vars
+    def step(self, var):
+        var.update = [u.clone() for u in var.get_update()]
+        return var
 
 class Zeros(Module):
     def __init__(self):
         super().__init__({})
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [torch.zeros_like(p) for p in vars.params]
-        return vars
+    def step(self, var):
+        var.update = [torch.zeros_like(p) for p in var.params]
+        return var
 
 class Ones(Module):
     def __init__(self):
         super().__init__({})
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [torch.ones_like(p) for p in vars.params]
-        return vars
+    def step(self, var):
+        var.update = [torch.ones_like(p) for p in var.params]
+        return var
 
 class Fill(Module):
     def __init__(self, value: float):
@@ -57,9 +57,9 @@ class Fill(Module):
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [torch.full_like(p, self.settings[p]['value']) for p in vars.params]
-        return vars
+    def step(self, var):
+        var.update = [torch.full_like(p, self.settings[p]['value']) for p in var.params]
+        return var
 
 class RandomSample(Module):
     def __init__(self, eps: float = 1, distribution: Distributions = 'normal'):
@@ -67,20 +67,20 @@ class RandomSample(Module):
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, vars):
-        vars.update = TensorList(vars.params).sample_like(
-            eps=self.get_settings('eps',params=vars.params), distribution=self.settings[vars.params[0]]['distribution']
+    def step(self, var):
+        var.update = TensorList(var.params).sample_like(
+            eps=[self.settings[p]['eps'] for p in var.params], distribution=self.settings[var.params[0]]['distribution']
         )
-        return vars
+        return var
 
 class Randn(Module):
     def __init__(self):
         super().__init__({})
 
     @torch.no_grad
-    def step(self, vars):
-        vars.update = [torch.randn_like(p) for p in vars.params]
-        return vars
+    def step(self, var):
+        var.update = [torch.randn_like(p) for p in var.params]
+        return var
 
 class Uniform(Module):
     def __init__(self, low: float, high: float):
@@ -88,25 +88,25 @@ class Uniform(Module):
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, vars):
-        low,high = self.get_settings('low','high', params=vars.params)
-        vars.update = [torch.empty_like(t).uniform_(l,h) for t,l,h in zip(vars.params, low, high)]
-        return vars
+    def step(self, var):
+        low,high = self.get_settings(var.params, 'low','high')
+        var.update = [torch.empty_like(t).uniform_(l,h) for t,l,h in zip(var.params, low, high)]
+        return var
 
 class GradToNone(Module):
     def __init__(self): super().__init__()
-    def step(self, vars):
-        vars.grad = None
-        return vars
+    def step(self, var):
+        var.grad = None
+        return var
 
 class UpdateToNone(Module):
     def __init__(self): super().__init__()
-    def step(self, vars):
-        vars.update = None
-        return vars
+    def step(self, var):
+        var.update = None
+        return var
 
 class Identity(Module):
     def __init__(self, *args, **kwargs): super().__init__()
-    def step(self, vars): return vars
+    def step(self, var): return var
 
 NoOp = Identity

torchzero/modules/optimizers/adagrad.py

@@ -1,18 +1,17 @@
 from operator import itemgetter
+from typing import Literal
 
 import torch
-
 from ...core import (
     Chainable,
     Module,
-    Preconditioner,
     Target,
-    TensorwisePreconditioner,
+    TensorwiseTransform,
     Transform,
-    Vars,
-    apply,
+    Var,
+    apply_transform,
 )
-from ...utils import NumberList, TensorList
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 from ...utils.linalg import matrix_power_eigh
 from ..functional import add_power_, lerp_power_, root
 
@@ -31,7 +30,6 @@ def adagrad_(
     inner: Module | None = None,
     params: list[torch.Tensor] | None = None,
     grads: list[torch.Tensor] | None = None,
-    vars: Vars | None = None,
 ):
     """returns `tensors_`"""
     clr = alpha / (1 + step * lr_decay)
@@ -40,7 +38,7 @@ def adagrad_(
 
     if inner is not None:
         assert params is not None
-        tensors_ = TensorList(apply(inner, tensors_, params=params, grads=grads, vars=vars))
+        tensors_ = TensorList(apply_transform(inner, tensors_, params=params, grads=grads))
 
     if use_sqrt: tensors_.div_(root(sq_sum_, p=pow, inplace=False).add_(eps)).mul_(clr)
     else: tensors_.div_(sq_sum_.add(eps)).mul_(clr)
@@ -79,19 +77,19 @@ class Adagrad(Transform):
             self.set_child('inner', inner)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         tensors = TensorList(tensors)
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        lr_decay,alpha,eps = self.get_settings('lr_decay', 'alpha', 'eps', params=params, cls=NumberList)
+        lr_decay,alpha,eps = unpack_dicts(settings, 'lr_decay', 'alpha', 'eps', cls=NumberList)
 
-        pow, use_sqrt = itemgetter('pow', 'use_sqrt')(self.settings[params[0]])
+        pow, use_sqrt = itemgetter('pow', 'use_sqrt')(settings[0])
 
-        sq_sum = self.get_state('sq_sum', params=params, cls=TensorList)
+        sq_sum = unpack_states(states, tensors, 'sq_sum', cls=TensorList)
 
         # initialize accumulator on 1st step
         if step == 1:
-            sq_sum.set_(tensors.full_like(self.get_settings('initial_accumulator_value', params=params)))
+            sq_sum.set_(tensors.full_like([s['initial_accumulator_value'] for s in settings]))
 
         return adagrad_(
             tensors,
@@ -107,40 +105,51 @@ class Adagrad(Transform):
             inner=self.children.get("inner", None),
             params=params,
             grads=grads,
-            vars=vars,
         )
 
 
 
-class FullMatrixAdagrad(TensorwisePreconditioner):
-    def __init__(self, beta: float | None = None, decay: float | None = None, concat_params=False, update_freq=1, inner: Chainable | None = None):
-        defaults = dict(beta=beta, decay=decay)
+class FullMatrixAdagrad(TensorwiseTransform):
+    def __init__(self, beta: float | None = None, decay: float | None = None, sqrt:bool=True, concat_params=False, update_freq=1, init: Literal['identity', 'zeros', 'ones', 'GGT'] = 'identity', inner: Chainable | None = None):
+        defaults = dict(beta=beta, decay=decay, sqrt=sqrt, init=init)
         super().__init__(defaults, uses_grad=False, concat_params=concat_params, update_freq=update_freq, inner=inner)
 
     @torch.no_grad
-    def update_tensor(self, tensor, param, grad, state, settings):
+    def update_tensor(self, tensor, param, grad, loss, state, settings):
         G = tensor.ravel()
         GG = torch.outer(G, G)
         decay = settings['decay']
         beta = settings['beta']
-
-        if 'GG' not in state: state['GG'] = torch.eye(GG.size(0), device=GG.device, dtype=GG.dtype)
+        init = settings['init']
+
+        if 'GG' not in state:
+            if init == 'identity': state['GG'] = torch.eye(GG.size(0), device=GG.device, dtype=GG.dtype)
+            elif init == 'zeros': state['GG'] =  torch.zeros_like(GG)
+            elif init == 'ones': state['GG'] = torch.ones_like(GG)
+            elif init == 'GGT': state['GG'] = GG.clone()
+            else: raise ValueError(init)
         if decay is not None: state['GG'].mul_(decay)
 
         if beta is not None: state['GG'].lerp_(GG, 1-beta)
         else: state['GG'].add_(GG)
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, state, settings):
+    def apply_tensor(self, tensor, param, grad, loss, state, settings):
         GG = state['GG']
+        sqrt = settings['sqrt']
 
         if tensor.numel() == 1:
-            return tensor / (GG**(1/2)).squeeze()
+            GG = GG.squeeze()
+            if sqrt: return tensor / GG.sqrt()
+            return tensor / GG
 
         try:
-            B = matrix_power_eigh(GG, -1/2)
+            if sqrt: B = matrix_power_eigh(GG, -1/2)
+            else: return torch.linalg.solve(GG, tensor.ravel()).view_as(tensor) # pylint:disable = not-callable
+
         except torch.linalg.LinAlgError:
-            return tensor.div_(tensor.abs().max()) # conservative scaling
+            scale = 1 / tensor.abs().max()
+            return tensor.mul_(scale.clip(min=torch.finfo(tensor.dtype).eps, max=1)) # conservative scaling
 
         return (B @ tensor.ravel()).view_as(tensor)