use factory_kwargs

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/compressed_tensors/transform/factory/hadamard.py

@@ -55,7 +55,8 @@ def create_transform(self, module: Module, args: TransformArgs):
         device = get_offloaded_device(module)
         exec_device = get_execution_device(module)
 
-        weight = self.weights.get(size, dtype, device, construct_device=exec_device)
+        factory_kwargs = {"construct_device": exec_device}
+        weight = self.weights.get(size, dtype, device, factory_kwargs=factory_kwargs)
         return HadamardTransform(weight, args)
 
     def _create_weight(

src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -62,6 +62,7 @@ def create_transform(self, module: Module, args: TransformArgs):
         return RandomMatrixTransform(weight, args)
 
     def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:
+        # TODO: verify that weight is invertable (has non-zero determinant)
         data = torch.rand(
             (size, size), generator=self.generator, dtype=dtype, device=device
         )

src/compressed_tensors/utils/helpers.py

@@ -15,10 +15,11 @@
 import contextlib
 import warnings
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional
 
 import numpy
 import torch
+from frozendict import frozendict
 from transformers import AutoConfig
 
 
@@ -373,16 +374,23 @@ class ParameterizedDefaultDict(dict):
 
     def __init__(self, default_factory: Callable[[Any], Any]):
         self.default_factory = default_factory
-        self._kwargs = {}
+        self._factory_kwargs = frozendict()
 
     def __missing__(self, key: Any) -> Any:
         if isinstance(key, tuple):
-            value = self.default_factory(*key, **self._kwargs)
+            value = self.default_factory(*key, **self._factory_kwargs)
         else:
-            value = self.default_factory(key, **self._kwargs)
+            value = self.default_factory(key, **self._factory_kwargs)
         self[key] = value
         return value
 
-    def get(self, *args, **kwargs) -> Any:
-        with patch_attr(self, "_kwargs", kwargs):
+    def get(self, *args, factory_kwargs: Mapping = frozendict()) -> Any:
+        """
+        Similar to `__getitem__`, but allows passing kwargs to factory function
+
+        :param \\*args: args whose tuple will value will be treated as key
+        :param factory_kwargs: keyword arguments to pass to `default_factory`
+        :return: dictionary entry for given key
+        """
+        with patch_attr(self, "_factory_kwargs", factory_kwargs):
             return self[args]