mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-16 22:46:38 +08:00
496 lines
21 KiB
Python
496 lines
21 KiB
Python
import unittest
|
|
from unittest import mock
|
|
import torch
|
|
import sys
|
|
import os
|
|
import json
|
|
|
|
# Add comfy to path
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
|
|
def has_gpu():
|
|
return torch.cuda.is_available()
|
|
|
|
from comfy.cli_args import args
|
|
if not has_gpu():
|
|
args.cpu = True
|
|
|
|
from comfy import ops
|
|
from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, TensorCoreNVFP4Layout
|
|
import comfy.utils
|
|
|
|
|
|
class SimpleModel(torch.nn.Module):
|
|
def __init__(self, operations=ops.disable_weight_init):
|
|
super().__init__()
|
|
self.layer1 = operations.Linear(10, 20, device="cpu", dtype=torch.bfloat16)
|
|
self.layer2 = operations.Linear(20, 30, device="cpu", dtype=torch.bfloat16)
|
|
self.layer3 = operations.Linear(30, 40, device="cpu", dtype=torch.bfloat16)
|
|
|
|
def forward(self, x):
|
|
x = self.layer1(x)
|
|
x = torch.nn.functional.relu(x)
|
|
x = self.layer2(x)
|
|
x = torch.nn.functional.relu(x)
|
|
x = self.layer3(x)
|
|
return x
|
|
|
|
|
|
class TestMixedPrecisionOps(unittest.TestCase):
|
|
|
|
@staticmethod
|
|
def _nvfp4_tensor(rows, columns, fill, scale=1.0):
|
|
params = TensorCoreNVFP4Layout.Params(
|
|
scale=None if scale is None else torch.tensor(scale),
|
|
orig_dtype=torch.float32,
|
|
orig_shape=(rows, columns),
|
|
block_scale=torch.ones(1),
|
|
)
|
|
return QuantizedTensor(
|
|
torch.full((rows, columns // 2), fill, dtype=torch.uint8),
|
|
"TensorCoreNVFP4Layout",
|
|
params,
|
|
)
|
|
|
|
@staticmethod
|
|
def _nvfp4_linear(in_features=16, out_features=8):
|
|
layer = ops.mixed_precision_ops({}).Linear(in_features, out_features, device="cpu", dtype=torch.float32)
|
|
layer.layout_type = "TensorCoreNVFP4Layout"
|
|
layer.quant_format = "nvfp4"
|
|
return layer
|
|
|
|
def test_all_layers_standard(self):
|
|
"""Test that model with no quantization works normally"""
|
|
# Create model
|
|
model = SimpleModel(operations=ops.mixed_precision_ops({}))
|
|
|
|
# Initialize weights manually
|
|
model.layer1.weight = torch.nn.Parameter(torch.randn(20, 10, dtype=torch.bfloat16))
|
|
model.layer1.bias = torch.nn.Parameter(torch.randn(20, dtype=torch.bfloat16))
|
|
model.layer2.weight = torch.nn.Parameter(torch.randn(30, 20, dtype=torch.bfloat16))
|
|
model.layer2.bias = torch.nn.Parameter(torch.randn(30, dtype=torch.bfloat16))
|
|
model.layer3.weight = torch.nn.Parameter(torch.randn(40, 30, dtype=torch.bfloat16))
|
|
model.layer3.bias = torch.nn.Parameter(torch.randn(40, dtype=torch.bfloat16))
|
|
|
|
# Initialize weight_function and bias_function
|
|
for layer in [model.layer1, model.layer2, model.layer3]:
|
|
layer.weight_function = []
|
|
layer.bias_function = []
|
|
|
|
# Forward pass
|
|
input_tensor = torch.randn(5, 10, dtype=torch.bfloat16)
|
|
output = model(input_tensor)
|
|
|
|
self.assertEqual(output.shape, (5, 40))
|
|
self.assertEqual(output.dtype, torch.bfloat16)
|
|
|
|
def test_mixed_precision_load(self):
|
|
"""Test loading a mixed precision model from state dict"""
|
|
# Configure mixed precision: layer1 is FP8, layer2 and layer3 are standard
|
|
layer_quant_config = {
|
|
"layer1": {
|
|
"format": "float8_e4m3fn",
|
|
"params": {}
|
|
},
|
|
"layer3": {
|
|
"format": "float8_e4m3fn",
|
|
"params": {}
|
|
}
|
|
}
|
|
|
|
# Create state dict with mixed precision
|
|
fp8_weight1 = torch.randn(20, 10, dtype=torch.float32).to(torch.float8_e4m3fn)
|
|
fp8_weight3 = torch.randn(40, 30, dtype=torch.float32).to(torch.float8_e4m3fn)
|
|
|
|
state_dict = {
|
|
# Layer 1: FP8 E4M3FN
|
|
"layer1.weight": fp8_weight1,
|
|
"layer1.bias": torch.randn(20, dtype=torch.bfloat16),
|
|
"layer1.weight_scale": torch.tensor(2.0, dtype=torch.float32),
|
|
|
|
# Layer 2: Standard BF16
|
|
"layer2.weight": torch.randn(30, 20, dtype=torch.bfloat16),
|
|
"layer2.bias": torch.randn(30, dtype=torch.bfloat16),
|
|
|
|
# Layer 3: FP8 E4M3FN
|
|
"layer3.weight": fp8_weight3,
|
|
"layer3.bias": torch.randn(40, dtype=torch.bfloat16),
|
|
"layer3.weight_scale": torch.tensor(1.5, dtype=torch.float32),
|
|
}
|
|
|
|
state_dict, _ = comfy.utils.convert_old_quants(state_dict, metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})})
|
|
# Create model and load state dict (strict=False because custom loading pops keys)
|
|
model = SimpleModel(operations=ops.mixed_precision_ops({}))
|
|
model.load_state_dict(state_dict, strict=False)
|
|
|
|
# Verify weights are wrapped in QuantizedTensor
|
|
self.assertIsInstance(model.layer1.weight, QuantizedTensor)
|
|
self.assertEqual(model.layer1.weight._layout_cls, "TensorCoreFP8E4M3Layout")
|
|
|
|
# Layer 2 should NOT be quantized
|
|
self.assertNotIsInstance(model.layer2.weight, QuantizedTensor)
|
|
|
|
# Layer 3 should be quantized
|
|
self.assertIsInstance(model.layer3.weight, QuantizedTensor)
|
|
self.assertEqual(model.layer3.weight._layout_cls, "TensorCoreFP8E4M3Layout")
|
|
|
|
# Verify scales were loaded
|
|
self.assertEqual(model.layer1.weight._params.scale.item(), 2.0)
|
|
self.assertEqual(model.layer3.weight._params.scale.item(), 1.5)
|
|
|
|
# Forward pass
|
|
input_tensor = torch.randn(5, 10, dtype=torch.bfloat16)
|
|
with torch.inference_mode():
|
|
output = model(input_tensor)
|
|
|
|
self.assertEqual(output.shape, (5, 40))
|
|
|
|
def test_state_dict_quantized_preserved(self):
|
|
"""Test that quantized weights are preserved in state_dict()"""
|
|
# Configure mixed precision
|
|
layer_quant_config = {
|
|
"layer1": {
|
|
"format": "float8_e4m3fn",
|
|
"params": {}
|
|
}
|
|
}
|
|
|
|
# Create and load model
|
|
fp8_weight = torch.randn(20, 10, dtype=torch.float32).to(torch.float8_e4m3fn)
|
|
state_dict1 = {
|
|
"layer1.weight": fp8_weight,
|
|
"layer1.bias": torch.randn(20, dtype=torch.bfloat16),
|
|
"layer1.weight_scale": torch.tensor(3.0, dtype=torch.float32),
|
|
"layer2.weight": torch.randn(30, 20, dtype=torch.bfloat16),
|
|
"layer2.bias": torch.randn(30, dtype=torch.bfloat16),
|
|
"layer3.weight": torch.randn(40, 30, dtype=torch.bfloat16),
|
|
"layer3.bias": torch.randn(40, dtype=torch.bfloat16),
|
|
}
|
|
|
|
state_dict1, _ = comfy.utils.convert_old_quants(state_dict1, metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})})
|
|
model = SimpleModel(operations=ops.mixed_precision_ops({}))
|
|
model.load_state_dict(state_dict1, strict=False)
|
|
|
|
# Save state dict
|
|
state_dict2 = model.state_dict()
|
|
|
|
# Verify layer1.weight is a QuantizedTensor with scale preserved
|
|
self.assertTrue(torch.equal(state_dict2["layer1.weight"].view(torch.uint8), fp8_weight.view(torch.uint8)))
|
|
self.assertEqual(state_dict2["layer1.weight_scale"].item(), 3.0)
|
|
self.assertEqual(model.layer1.weight._layout_cls, "TensorCoreFP8E4M3Layout")
|
|
|
|
# Verify non-quantized layers are standard tensors
|
|
self.assertNotIsInstance(state_dict2["layer2.weight"], QuantizedTensor)
|
|
self.assertNotIsInstance(state_dict2["layer3.weight"], QuantizedTensor)
|
|
|
|
def test_weight_function_compatibility(self):
|
|
"""Test that weight_function (LoRA) works with quantized layers"""
|
|
# Configure FP8 quantization
|
|
layer_quant_config = {
|
|
"layer1": {
|
|
"format": "float8_e4m3fn",
|
|
"params": {}
|
|
}
|
|
}
|
|
|
|
# Create and load model
|
|
fp8_weight = torch.randn(20, 10, dtype=torch.float32).to(torch.float8_e4m3fn)
|
|
state_dict = {
|
|
"layer1.weight": fp8_weight,
|
|
"layer1.bias": torch.randn(20, dtype=torch.bfloat16),
|
|
"layer1.weight_scale": torch.tensor(2.0, dtype=torch.float32),
|
|
"layer2.weight": torch.randn(30, 20, dtype=torch.bfloat16),
|
|
"layer2.bias": torch.randn(30, dtype=torch.bfloat16),
|
|
"layer3.weight": torch.randn(40, 30, dtype=torch.bfloat16),
|
|
"layer3.bias": torch.randn(40, dtype=torch.bfloat16),
|
|
}
|
|
|
|
state_dict, _ = comfy.utils.convert_old_quants(state_dict, metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})})
|
|
model = SimpleModel(operations=ops.mixed_precision_ops({}))
|
|
model.load_state_dict(state_dict, strict=False)
|
|
|
|
# Add a weight function (simulating LoRA)
|
|
# This should trigger dequantization during forward pass
|
|
def apply_lora(weight):
|
|
lora_delta = torch.randn_like(weight) * 0.01
|
|
return weight + lora_delta
|
|
|
|
model.layer1.weight_function.append(apply_lora)
|
|
|
|
# Forward pass should work with LoRA (triggers weight_function path)
|
|
input_tensor = torch.randn(5, 10, dtype=torch.bfloat16)
|
|
output = model(input_tensor)
|
|
|
|
self.assertEqual(output.shape, (5, 40))
|
|
|
|
def test_error_handling_unknown_format(self):
|
|
"""Test that unknown formats raise error"""
|
|
# Configure with unknown format
|
|
layer_quant_config = {
|
|
"layer1": {
|
|
"format": "unknown_format_xyz",
|
|
"params": {}
|
|
}
|
|
}
|
|
|
|
# Create state dict
|
|
state_dict = {
|
|
"layer1.weight": torch.randn(20, 10, dtype=torch.bfloat16),
|
|
"layer1.bias": torch.randn(20, dtype=torch.bfloat16),
|
|
"layer2.weight": torch.randn(30, 20, dtype=torch.bfloat16),
|
|
"layer2.bias": torch.randn(30, dtype=torch.bfloat16),
|
|
"layer3.weight": torch.randn(40, 30, dtype=torch.bfloat16),
|
|
"layer3.bias": torch.randn(40, dtype=torch.bfloat16),
|
|
}
|
|
|
|
state_dict, _ = comfy.utils.convert_old_quants(state_dict, metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})})
|
|
|
|
# Load should raise KeyError for unknown format in QUANT_FORMAT_MIXINS
|
|
model = SimpleModel(operations=ops.mixed_precision_ops({}))
|
|
with self.assertRaises(KeyError):
|
|
model.load_state_dict(state_dict, strict=False)
|
|
|
|
def test_int8_convrot_metadata_loads_into_params(self):
|
|
"""ConvRot metadata must reach TensorWiseINT8Layout params."""
|
|
torch.manual_seed(123)
|
|
layer_quant_config = {
|
|
"layer": {
|
|
"format": "int8_tensorwise",
|
|
"convrot": True,
|
|
"convrot_groupsize": 256,
|
|
}
|
|
}
|
|
weight = torch.randn(16, 256, dtype=torch.bfloat16)
|
|
bias = torch.randn(16, dtype=torch.bfloat16)
|
|
q_weight = QuantizedTensor.from_float(
|
|
weight,
|
|
"TensorWiseINT8Layout",
|
|
per_channel=True,
|
|
convrot=True,
|
|
convrot_groupsize=256,
|
|
)
|
|
state_dict = {
|
|
"layer.weight": q_weight._qdata,
|
|
"layer.bias": bias,
|
|
"layer.weight_scale": q_weight._params.scale,
|
|
}
|
|
|
|
state_dict, _ = comfy.utils.convert_old_quants(
|
|
state_dict,
|
|
metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})},
|
|
)
|
|
model = torch.nn.Module()
|
|
model.layer = ops.mixed_precision_ops({}).Linear(256, 16, device="cpu", dtype=torch.bfloat16)
|
|
model.load_state_dict(state_dict, strict=False)
|
|
|
|
self.assertIsInstance(model.layer.weight, QuantizedTensor)
|
|
self.assertEqual(model.layer.weight._layout_cls, "TensorWiseINT8Layout")
|
|
self.assertTrue(model.layer.weight._params.convrot)
|
|
self.assertEqual(model.layer.weight._params.convrot_groupsize, 256)
|
|
|
|
input_tensor = torch.randn(4, 256, dtype=torch.bfloat16)
|
|
loaded_out = model.layer(input_tensor)
|
|
ref_out = torch.nn.functional.linear(input_tensor, q_weight, bias)
|
|
self.assertTrue(torch.equal(loaded_out, ref_out))
|
|
|
|
fp16_input = input_tensor.to(torch.float16)
|
|
loaded_fp16_out = model.layer(fp16_input)
|
|
ref_fp16_out = torch.nn.functional.linear(
|
|
fp16_input,
|
|
q_weight.to(dtype=torch.float16),
|
|
bias.to(dtype=torch.float16),
|
|
)
|
|
self.assertTrue(torch.equal(loaded_fp16_out, ref_fp16_out))
|
|
|
|
saved = model.state_dict()
|
|
saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes())
|
|
self.assertTrue(saved_conf["convrot"])
|
|
|
|
def test_convrot_w4a4_loads_into_params(self):
|
|
"""ConvRot W4A4 checkpoints must load as the dedicated kitchen layout."""
|
|
if "convrot_w4a4" not in QUANT_ALGOS:
|
|
self.skipTest("comfy_kitchen does not provide ConvRot W4A4")
|
|
|
|
torch.manual_seed(456)
|
|
layer_quant_config = {
|
|
"layer": {
|
|
"format": "convrot_w4a4",
|
|
"convrot_groupsize": 256,
|
|
"linear_dtype": "int8",
|
|
}
|
|
}
|
|
weight = torch.randn(16, 256, dtype=torch.bfloat16)
|
|
bias = torch.randn(16, dtype=torch.bfloat16)
|
|
q_weight = QuantizedTensor.from_float(
|
|
weight,
|
|
"TensorCoreConvRotW4A4Layout",
|
|
convrot_groupsize=256,
|
|
quant_group_size=64,
|
|
)
|
|
state_dict = {
|
|
"layer.weight": q_weight._qdata,
|
|
"layer.bias": bias,
|
|
"layer.weight_scale": q_weight._params.scale,
|
|
}
|
|
|
|
state_dict, _ = comfy.utils.convert_old_quants(
|
|
state_dict,
|
|
metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})},
|
|
)
|
|
model = torch.nn.Module()
|
|
model.layer = ops.mixed_precision_ops({}).Linear(256, 16, device="cpu", dtype=torch.bfloat16)
|
|
model.load_state_dict(state_dict, strict=False)
|
|
|
|
self.assertIsInstance(model.layer.weight, QuantizedTensor)
|
|
self.assertEqual(model.layer.weight._layout_cls, "TensorCoreConvRotW4A4Layout")
|
|
self.assertEqual(model.layer.weight._params.convrot_groupsize, 256)
|
|
self.assertEqual(model.layer.weight._params.quant_group_size, 64)
|
|
self.assertEqual(model.layer.weight._params.linear_dtype, "int8")
|
|
|
|
input_tensor = torch.randn(4, 256, dtype=torch.bfloat16)
|
|
loaded_out = model.layer(input_tensor)
|
|
ref_out = torch.nn.functional.linear(input_tensor, q_weight, bias)
|
|
self.assertTrue(torch.equal(loaded_out, ref_out))
|
|
|
|
saved = model.state_dict()
|
|
saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes())
|
|
self.assertEqual(saved_conf["format"], "convrot_w4a4")
|
|
self.assertEqual(saved_conf["convrot_groupsize"], 256)
|
|
self.assertEqual(saved_conf["linear_dtype"], "int8")
|
|
self.assertNotIn("quant_group_size", saved_conf)
|
|
|
|
def test_nvfp4_quantized_mm_dispatch_with_bias_and_rank_restore(self):
|
|
input_tensor = torch.randn(2, 3, 16)
|
|
quantized_input = self._nvfp4_tensor(6, 16, 1)
|
|
quantized_weight = self._nvfp4_tensor(8, 16, 2)
|
|
bias = torch.arange(8, dtype=torch.float32)
|
|
kernel_output = torch.arange(48, dtype=torch.float32).reshape(6, 8)
|
|
layer = self._nvfp4_linear()
|
|
layer.weight = quantized_weight
|
|
|
|
with (
|
|
mock.patch.object(QuantizedTensor, "from_float", return_value=quantized_input) as quantize,
|
|
mock.patch("comfy_kitchen.scaled_mm_nvfp4", return_value=kernel_output) as scaled_mm,
|
|
mock.patch.object(QuantizedTensor, "dequantize", autospec=True) as dequantize,
|
|
mock.patch.object(ops, "cast_bias_weight", return_value=(quantized_weight, bias, None)),
|
|
mock.patch.object(ops, "uncast_bias_weight"),
|
|
):
|
|
output = layer(input_tensor)
|
|
|
|
self.assertTrue(torch.equal(output, (kernel_output + bias).reshape(2, 3, 8)))
|
|
self.assertEqual(quantize.call_args.args[0].shape, (6, 16))
|
|
self.assertEqual(scaled_mm.call_count, 1)
|
|
self.assertEqual(dequantize.call_count, 0)
|
|
|
|
def test_nvfp4_quantized_mm_dispatch_without_bias(self):
|
|
quantized_input = self._nvfp4_tensor(5, 16, 1)
|
|
quantized_weight = self._nvfp4_tensor(8, 16, 2)
|
|
kernel_output = torch.randn(5, 8)
|
|
layer = self._nvfp4_linear()
|
|
|
|
with (
|
|
mock.patch("comfy_kitchen.scaled_mm_nvfp4", return_value=kernel_output) as scaled_mm,
|
|
mock.patch.object(QuantizedTensor, "dequantize", autospec=True) as dequantize,
|
|
):
|
|
output = layer._forward(quantized_input, quantized_weight, None)
|
|
|
|
self.assertIs(output, kernel_output)
|
|
self.assertEqual(scaled_mm.call_count, 1)
|
|
self.assertEqual(dequantize.call_count, 0)
|
|
|
|
def test_nvfp4_mm_backend_fallbacks_dequantize(self):
|
|
input_tensor = torch.randn(5, 16)
|
|
weight = torch.randn(8, 16)
|
|
bias = torch.randn(8)
|
|
expected = torch.nn.functional.linear(input_tensor, weight, bias)
|
|
quantized_input = self._nvfp4_tensor(5, 16, 1)
|
|
quantized_weight = self._nvfp4_tensor(8, 16, 2)
|
|
dequantized = {
|
|
(5, 16): input_tensor,
|
|
(8, 16): weight,
|
|
}
|
|
layer = self._nvfp4_linear()
|
|
|
|
def dequantize(tensor):
|
|
value = dequantized[tensor._params.orig_shape if not tensor._params.transposed else tuple(reversed(tensor._params.orig_shape))]
|
|
return value.t() if tensor._params.transposed else value
|
|
|
|
for error in (RuntimeError("no supported backend"), TypeError("unsupported kernel signature")):
|
|
with self.subTest(error=type(error).__name__):
|
|
with (
|
|
mock.patch("comfy_kitchen.scaled_mm_nvfp4", side_effect=error) as scaled_mm,
|
|
mock.patch.object(QuantizedTensor, "dequantize", autospec=True, side_effect=dequantize) as dequantize_call,
|
|
):
|
|
output = layer._forward(quantized_input, quantized_weight, bias)
|
|
|
|
self.assertTrue(torch.equal(output, expected))
|
|
self.assertEqual(scaled_mm.call_count, 1)
|
|
self.assertEqual(dequantize_call.call_count, 2)
|
|
|
|
def test_nvfp4_mm_unexpected_failures_remain_loud(self):
|
|
class Cancellation(BaseException):
|
|
pass
|
|
|
|
quantized_input = self._nvfp4_tensor(5, 16, 1)
|
|
quantized_weight = self._nvfp4_tensor(8, 16, 2)
|
|
layer = self._nvfp4_linear()
|
|
|
|
for error in (ValueError("invalid metadata"), Cancellation("cancelled")):
|
|
with self.subTest(error=type(error).__name__):
|
|
with (
|
|
mock.patch("comfy_kitchen.scaled_mm_nvfp4", side_effect=error),
|
|
mock.patch.object(QuantizedTensor, "dequantize", autospec=True) as dequantize,
|
|
):
|
|
with self.assertRaises(type(error)):
|
|
layer._forward(quantized_input, quantized_weight, None)
|
|
self.assertEqual(dequantize.call_count, 0)
|
|
|
|
def test_nvfp4_mm_invalid_quantization_metadata_remains_loud(self):
|
|
quantized_input = self._nvfp4_tensor(5, 16, 1)
|
|
layer = self._nvfp4_linear()
|
|
|
|
missing_params = self._nvfp4_tensor(8, 16, 2)
|
|
del missing_params._params
|
|
with self.assertRaises(AttributeError):
|
|
layer._forward(quantized_input, missing_params, None)
|
|
|
|
unknown_layout = self._nvfp4_tensor(8, 16, 2)
|
|
unknown_layout._layout_cls = "UnknownLayout"
|
|
with self.assertRaises(KeyError):
|
|
layer._forward(quantized_input, unknown_layout, None)
|
|
|
|
missing_scale = self._nvfp4_tensor(8, 16, 2, scale=None)
|
|
with mock.patch("comfy_kitchen.scaled_mm_nvfp4", side_effect=ValueError("missing scale")):
|
|
with self.assertRaisesRegex(ValueError, "missing scale"):
|
|
layer._forward(quantized_input, missing_scale, None)
|
|
|
|
def test_nvfp4_mm_guard_preserves_other_linear_paths(self):
|
|
quantized_input = self._nvfp4_tensor(5, 16, 1)
|
|
quantized_weight = self._nvfp4_tensor(8, 16, 2)
|
|
regular_input = torch.randn(5, 16)
|
|
regular_weight = torch.randn(8, 16)
|
|
bias = torch.randn(8)
|
|
sentinel = torch.randn(5, 8)
|
|
layer = self._nvfp4_linear()
|
|
|
|
cases = (
|
|
("TensorCoreFP8E4M3Layout", quantized_input, quantized_weight),
|
|
("TensorCoreNVFP4Layout", regular_input, quantized_weight),
|
|
("TensorCoreNVFP4Layout", quantized_input, regular_weight),
|
|
)
|
|
for layout_type, input_value, weight_value in cases:
|
|
with self.subTest(layout_type=layout_type, input_type=type(input_value).__name__, weight_type=type(weight_value).__name__):
|
|
with (
|
|
mock.patch.object(torch.nn.functional, "linear", return_value=sentinel) as functional_linear,
|
|
mock.patch.object(torch, "mm") as mm,
|
|
):
|
|
layer.layout_type = layout_type
|
|
output = layer._forward(input_value, weight_value, bias)
|
|
|
|
self.assertIs(output, sentinel)
|
|
functional_linear.assert_called_once_with(input_value, weight_value, bias)
|
|
mm.assert_not_called()
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|