narrow arguments

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/compressed_tensors/transform/factory/hadamard.py

@@ -40,7 +40,7 @@ def __init__(self, name: str, scheme: TransformScheme, seed: int = 42):
 
     def create_transform(self, module: Module, args: TransformArgs):
         assert isinstance(module, Linear)
-        size = get_matrix_size(module, args)
+        size = get_matrix_size(module, args.location)
         dtype = module.weight.dtype
         device = get_offloaded_device(module)
 
@@ -76,4 +76,4 @@ def forward(self, value: Tensor) -> Tensor:
         # if self.permutation is not None:
         #    weight = apply_permutation(weight, self.permutation)
 
-        return apply_transform_weight(weight, value, self.args)
+        return apply_transform_weight(weight, value, self.args.location)

src/compressed_tensors/transform/factory/matrix_multiply.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
-
 import torch
 from compressed_tensors.transform import TransformArgs, TransformScheme
 from compressed_tensors.transform.factory.base import TransformBase, TransformFactory
@@ -42,7 +40,7 @@ def __init__(
 
     def create_transform(self, module: Module, args: TransformArgs):
         assert isinstance(module, Linear)
-        size = get_matrix_size(module, args)
+        size = get_matrix_size(module, args.location)
         dtype = module.weight.dtype
         device = get_offloaded_device(module)
 
@@ -68,7 +66,7 @@ def __init__(self, weight: Tensor, args: TransformArgs):
         self.args = args
 
     def forward(self, value: Tensor) -> Parameter:
-        return apply_transform_weight(self.weight, value, self.args)
+        return apply_transform_weight(self.weight, value, self.args.location)
 
 
 def high_precision_invert(weight: Tensor) -> Tensor:

src/compressed_tensors/transform/utils/utils.py

@@ -12,18 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Literal
-
 import torch
-from compressed_tensors.transform import TransformArgs, TransformLocation
+from compressed_tensors.transform import TransformLocation
 
 
 __all__ = ["get_matrix_size", "apply_transform_weight", "apply_permutation"]
 
 
-def get_matrix_size(module: torch.nn.Module, args: TransformArgs) -> int:
+def get_matrix_size(module: torch.nn.Module, location: TransformLocation) -> int:
     assert isinstance(module, torch.nn.Linear)
-    if args.location in ("input", TransformLocation.WEIGHT_INPUT):
+    if location in ("input", TransformLocation.WEIGHT_INPUT):
         return module.in_features
     else:
         return module.out_features
@@ -32,7 +30,7 @@ def get_matrix_size(module: torch.nn.Module, args: TransformArgs) -> int:
 def apply_transform_weight(
     weight: torch.Tensor,
     value: torch.Tensor,
-    args: TransformArgs,  # TODO: only pass location
+    location: TransformLocation,
 ) -> torch.Tensor:
     # let x          be input activation
     #     W          be weight,
@@ -57,16 +55,16 @@ def apply_transform_weight(
     #             = y U
     #             = yh
 
-    if args.location == TransformLocation.INPUT:
+    if location == TransformLocation.INPUT:
         return value @ weight
 
-    elif args.location == TransformLocation.WEIGHT_INPUT:
+    elif location == TransformLocation.WEIGHT_INPUT:
         return value @ weight.T
 
-    elif args.location == TransformLocation.WEIGHT_OUTPUT:
+    elif location == TransformLocation.WEIGHT_OUTPUT:
         return weight.T @ value
 
-    elif args.location == TransformLocation.OUTPUT:
+    elif location == TransformLocation.OUTPUT:
         return value @ weight