Fix parsing einsum with a single input (#1311)

Author: vloncar

hls4ml/converters/pytorch/core.py

@@ -166,7 +166,10 @@ def parse_einsum_layer(operation, layer_name, input_names, input_shapes, node, c
 
     layer = {}
 
-    if len(input_names) != 2:
+    if len(input_names) == 1:
+        input_names += input_names
+        input_shapes += input_shapes
+    elif len(input_names) > 2:
         raise Exception('Only einsum operations with two inputs are supported')
     layer['class_name'] = 'Einsum'
     layer['name'] = layer_name

test/pytest/test_pytorch_api.py

@@ -895,6 +895,18 @@ def forward(self, x, y):
         return torch.einsum('bij,bjk->bik', x, y)
 
 
+class EinsumSingleInput(nn.Module):
+    def __init__(self, input_dim=8):
+        super().__init__()
+        self.input_dim = input_dim
+        self.linear = nn.Linear(self.input_dim, self.input_dim)
+
+    def forward(self, x):
+        """using torch einsum to get the dot product"""
+        out = self.linear(x)
+        return torch.einsum("ij,ij->i", out, out)
+
+
 @pytest.mark.parametrize('backend', ['Vivado', 'Vitis'])
 @pytest.mark.parametrize('io_type', ['io_parallel'])
 def test_einsum_outer_product(backend, io_type):
@@ -953,3 +965,32 @@ def test_einsum_batch_matmul(backend, io_type):
     hls_prediction = np.reshape(hls_model.predict([X_input, Y_input]), pytorch_prediction.shape)
 
     np.testing.assert_allclose(hls_prediction, pytorch_prediction, rtol=1e-2, atol=0.01)
+
+
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis'])
+@pytest.mark.parametrize('io_type', ['io_parallel'])
+def test_einsum_single_input(backend, io_type):
+
+    model = EinsumSingleInput()
+    model.eval()
+
+    X_input = np.random.rand(3, 8)
+
+    pytorch_prediction = model(torch.Tensor(X_input)).detach().numpy()
+
+    config = config_from_pytorch_model(
+        model,
+        [(None, 8)],
+        default_precision='ap_fixed<16,6>',
+        channels_last_conversion="internal",
+        transpose_outputs=False,
+    )
+    output_dir = str(test_root_path / f'hls4mlprj_pytorch_einsum_single_input_{backend}_{io_type}')
+
+    hls_model = convert_from_pytorch_model(model, hls_config=config, output_dir=output_dir, backend=backend, io_type=io_type)
+
+    hls_model.compile()
+
+    hls_prediction = np.reshape(hls_model.predict(X_input), pytorch_prediction.shape)
+
+    np.testing.assert_allclose(hls_prediction, pytorch_prediction, rtol=1e-2, atol=0.01)