layout_tensor -> _layout

Author: jainapurva

test/dtypes/test_affine_quantized.py

@@ -31,7 +31,7 @@ def get_quantization_functions(do_sparse: bool, do_int4: bool):
         base_functions.append(int4_weight_only(group_size=32))
 
     if do_sparse:
-        base_functions.append(int8_dynamic_activation_int8_weight(layout_type=SemiSparseLayout()))
+        base_functions.append(int8_dynamic_activation_int8_weight(_layout=SemiSparseLayout()))
 
     if is_cuda_8_9:
         base_functions.append(float8_weight_only())

test/dtypes/test_floatx.py

@@ -81,8 +81,8 @@ def test_to_copy_device(self, ebits, mbits):
         x = torch.randn(256, 64)
         scale = choose_qparams_affine_floatx(x, ebits, mbits)
         x = quantize_affine_floatx(x, scale, ebits, mbits)
-        layout_type = FloatxTensorCoreLayout(ebits, mbits)
-        floatx_tensor_impl = FloatxTensorCoreAQTTensorImpl.from_plain(x, scale, None, layout_type).cuda()
+        _layout = FloatxTensorCoreLayout(ebits, mbits)
+        floatx_tensor_impl = FloatxTensorCoreAQTTensorImpl.from_plain(x, scale, None, _layout).cuda()
         assert floatx_tensor_impl.device.type == "cuda"
         floatx_tensor_impl = floatx_tensor_impl.cpu()
         assert floatx_tensor_impl.device.type == "cpu"

test/integration/test_integration.py

@@ -876,7 +876,7 @@ def test_int4_weight_only_quant_subclass_api_grouped(self, device, dtype):
         for test_shape in ([(256, 256, 16)] + ([(256, 256, 8)] if device=='cuda' else [])):
             for groupsize in [64, 32]:
                 for inner_k_tiles in [4, 2]:
-                    kwargs = {"groupsize": groupsize, "layout_type": TensorCoreTiledLayout(inner_k_tiles=inner_k_tiles)}
+                    kwargs = {"groupsize": groupsize, "_layout": TensorCoreTiledLayout(inner_k_tiles=inner_k_tiles)}
 
                     def api(mod):
                         kwargs_copy = kwargs.copy()
@@ -888,7 +888,7 @@ def api(mod):
                                 unwrap_tensor_subclass(mod)
                         else:
                             kwargs_copy["inner_k_tiles"] = inner_k_tiles
-                            del kwargs_copy["layout_type"]
+                            del kwargs_copy["_layout"]
                             change_linear_weights_to_int4_woqtensors(mod, **kwargs_copy)
 
                     self._test_lin_weight_subclass_api_impl(

test/sparsity/test_marlin.py

@@ -50,7 +50,7 @@ def test_quant_sparse_marlin_layout_eager(self):
         dense_result = model_copy(self.input.bfloat16()).half()
 
         # Sparse + quantized
-        quantize_(self.model, int4_weight_only(layout_type=MarlinSparseLayout()))
+        quantize_(self.model, int4_weight_only(_layout=MarlinSparseLayout()))
         sparse_result = self.model(self.input)
 
         assert torch.allclose(dense_result, sparse_result, atol=3e-1), "Results are not close"
@@ -67,7 +67,7 @@ def test_quant_sparse_marlin_layout_compile(self):
         dense_result = model_copy(self.input.bfloat16()).half()
 
         # Sparse + quantized
-        quantize_(self.model, int4_weight_only(layout_type=MarlinSparseLayout()))
+        quantize_(self.model, int4_weight_only(_layout=MarlinSparseLayout()))
         self.model.forward = torch.compile(self.model.forward, fullgraph=True)
         sparse_result = self.model(self.input)
 

test/sparsity/test_sparse_api.py

@@ -74,7 +74,7 @@ def test_quant_semi_sparse(self, compile):
 
         quantize_(
             model,
-            int8_dynamic_activation_int8_weight(layout_type=SemiSparseLayout()),
+            int8_dynamic_activation_int8_weight(_layout=SemiSparseLayout()),
         )
         if compile:
             model = torch.compile(model)
@@ -108,7 +108,7 @@ def test_sparse_marlin(self, compile):
         dense_result = model_copy(input.bfloat16()).half()
 
         # Sparse + quantized
-        quantize_(model, int4_weight_only(layout_type=MarlinSparseLayout()))
+        quantize_(model, int4_weight_only(_layout=MarlinSparseLayout()))
         if compile:
             model = torch.compile(model)
         sparse_result = model(input)
@@ -190,7 +190,7 @@ def test_sparse(self, compile):
         quantize_(
             model,
             int8_dynamic_activation_int8_weight(
-                layout_type=BlockSparseLayout(blocksize=64)
+                _layout=BlockSparseLayout(blocksize=64)
             ),
         )
         if compile:

torchao/_models/llama/eval.py

@@ -98,7 +98,7 @@ def run_evaluation(
             quantize_(model, uintx_weight_only(dtype, group_size, use_hqq=use_hqq))
         if "marlin" in quantization:
             from torchao.dtypes import MarlinSparseLayout
-            quantize_(model, int4_weight_only(layout_type=MarlinSparseLayout()))
+            quantize_(model, int4_weight_only(_layout=MarlinSparseLayout()))
         if "int4wo" in quantization and "gptq" in quantization:
             # avoid circular imports
             from torchao._models._eval import InputRecorder

torchao/_models/llama/generate.py

@@ -231,7 +231,7 @@ def main(
             quantize_(model, int4_weight_only(group_size=groupsize))
         if "marlin" in quantization:
             from torchao.dtypes import MarlinSparseLayout
-            quantize_(model, int4_weight_only(layout_type=MarlinSparseLayout()))
+            quantize_(model, int4_weight_only(_layout=MarlinSparseLayout()))
         if "fp6" in quantization:
             quantize_(model, fpx_weight_only(3, 2))
         if quantization.startswith("awq"):

torchao/_models/sam/eval_combo.py

@@ -315,7 +315,7 @@ def mlp_only(mod, name):
                   int8_dynamic_activation_int8_weight(),
                   attn_only)
         quantize_(predictor.model.image_encoder,
-                  int8_dynamic_activation_int8_weight(layout_type=SemiSparseLayout()),
+                  int8_dynamic_activation_int8_weight(_layout=SemiSparseLayout()),
                   mlp_lin1_only)
         sparsify_(predictor.model.image_encoder,
                   semi_sparse_weight(),
@@ -330,7 +330,7 @@ def mlp_only(mod, name):
         quantize_(predictor.model.image_encoder,
                   int8_dynamic_activation_int8_weight(),
                   attn_only)
-        quantize_(predictor.model.image_encoder, int4_weight_only(layout_type=MarlinSparseLayout()), mlp_lin1_only)
+        quantize_(predictor.model.image_encoder, int4_weight_only(_layout=MarlinSparseLayout()), mlp_lin1_only)
         sparsify_(predictor.model.image_encoder,
                   semi_sparse_weight(),
                   mlp_lin2_only)