fix(pytorch): place unsqueeze axis at the correct position

hls4ml/converters/pytorch/reshape.py

@@ -78,9 +78,16 @@ def parse_unsqueeze_layer(operation, layer_name, input_names, input_shapes, node
         squeeze_dim = node.args[1]
     else:  # Specified as unsqueeze(x, dim=n)
         squeeze_dim = node.kwargs['dim']
-    # insert() will add an element before the index, unsqueeze expects the location
-    index = output_shape.index(output_shape[squeeze_dim])  # + 1
-    output_shape.insert(index, 1)
+    # torch.unsqueeze inserts a new axis of size 1 at position 'squeeze_dim' and accepts
+    # dim in [-(D+1), D]. Reject out-of-range values (list.insert would otherwise silently
+    # clamp them to a wrong shape) and normalize negative dims, so the axis lands at the
+    # correct location regardless of duplicate dimension sizes.
+    ndim = len(output_shape)
+    if not -(ndim + 1) <= squeeze_dim <= ndim:
+        raise Exception(f'Dimension {squeeze_dim} is out of range for unsqueeze of a {ndim}D tensor')
+    if squeeze_dim < 0:
+        squeeze_dim += ndim + 1
+    output_shape.insert(squeeze_dim, 1)
 
     layer['target_shape'] = output_shape.copy()
     if layer['target_shape'][0] is None:

test/pytest/test_pytorch_api.py

@@ -673,6 +673,50 @@ def test_squeeze(test_case_id, backend, io_type):
         assert list(hls_model.get_layers())[3].attributes['target_shape'] == [3]
 
 
+class UnsqueezeModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(5, 4, bias=False)
+        nn.init.ones_(self.linear.weight)  # This test is not about precision, so put 1's here
+
+    def forward(self, x):
+        x = self.linear(x)  # (1, 5) -> (1, 4)
+        x = torch.unsqueeze(x, dim=-1)  # (1, 4) -> (1, 4, 1)
+        x = torch.relu(x)  # (1, 4, 1)
+        return x
+
+
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus', 'oneAPI'])
+@pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream'])
+def test_unsqueeze(test_case_id, backend, io_type):
+    # Regression test: torch.unsqueeze(x, dim=-1) must insert the size-1 axis as the *last*
+    # dimension. The previous parser located the new axis with list.index() on the dimension
+    # value, which placed it at the wrong position for negative dims (or whenever the indexed
+    # dimension shared its size with an earlier one).
+    model = UnsqueezeModel()
+    model.eval()
+
+    X_input = np.random.rand(1, 5)
+
+    pytorch_prediction = model(torch.Tensor(X_input)).detach().numpy().flatten()
+
+    config = config_from_pytorch_model(model, (5,))
+    del config['Model']['ChannelsLastConversion']  # We don't want anything touched for this test
+    output_dir = str(test_root_path / test_case_id)
+
+    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 = hls_model.predict(X_input).flatten()
+
+    np.testing.assert_allclose(hls_prediction, pytorch_prediction, rtol=1e-2, atol=0.01)
+
+    # The reshape (or its io_stream Repack counterpart) must report (4, 1), not (1, 4).
+    reshape_layer = next(layer for layer in hls_model.get_layers() if 'unsqueeze' in layer.name)
+    assert reshape_layer.attributes['target_shape'] == [4, 1]
+
+
 @pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'Quartus', 'oneAPI'])
 def test_flatten(test_case_id, backend):
     input = torch.randn(1, 1, 5, 5)