def forward(self,
tensor_list: Union[List[Tensor], List[QuantTensor]],
dim: int = 1) -> Union[Tensor, QuantTensor]:
quant_tensor_list = [self.unpack_input(t) for t in tensor_list]
# shortcut execution through the export impl during export
if self.export_mode:
return self.export_handler([qt.value for qt in quant_tensor_list])
quant_tensor_list = [self.input_quant(qt) for qt in quant_tensor_list]
# trigger an assert if scale factors and bit widths are None or different
output = QuantTensor.cat(quant_tensor_list, dim=dim)
quant_output = self.output_quant(output)
return self.pack_output(quant_output)
def cat_handler(*args, **kwargs):
from brevitas.quant_tensor import QuantTensor
return QuantTensor.cat(*args, **kwargs)