feat: improve error message when refolding with missing dims

foldedtensor/__init__.py

@@ -402,9 +402,16 @@ def refold(self, *dims: Union[Sequence[Union[int, str]], int, str]):
                 "sequence or each arguments to be ints or strings"
             )
             dims = dims[0]
-        dims = tuple(
-            dim if isinstance(dim, int) else self.full_names.index(dim) for dim in dims
-        )
+        try:
+            dims = tuple(
+                dim if isinstance(dim, int) else self.full_names.index(dim)
+                for dim in dims
+            )
+        except ValueError:
+            raise ValueError(
+                f"Folded tensor with available dimensions {self.full_names} "
+                f"could not be refolded with dimensions {list(dims)}"
+            )
 
         if dims == self.data_dims:
             return self

tests/test_folded_tensor.py

@@ -431,3 +431,18 @@ def test_hashable_lengths():
     assert tensor.lengths is embedding(tensor).lengths
     assert hash(tensor.lengths) is not None
     assert hash(tensor.lengths) == hash(embedding(tensor).lengths)
+
+
+def test_missing_dims():
+    tensor = as_folded_tensor(
+        [
+            [0, 1, 2],
+            [3, 4],
+        ],
+        full_names=("sample", "token"),
+        dtype=torch.long,
+    )
+    with pytest.raises(ValueError) as e:
+        tensor.refold("line", "token")
+
+    assert "line" in str(e.value)