Fix bug in BaseMCRTensor:__torch_function__()

There is a bug when args is a collection of collections instead of a
plain tuple. In this case, the old args parsing was unable to search for
block_size in nested structures containing BaseMCRTensor.

For instance, let `a` and `b` be BaseMCRTensor variables. Calling
`torch.stack((a, b))` results in error, as the `args` received in
`__torch_function__()` is a nested tuple.

Author: Zeldax64

torchhd/tensors/basemcr.py

@@ -359,7 +359,20 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
         if kwargs is None:
             kwargs = {}
 
-        block_sizes = set(a.block_size for a in args if hasattr(a, "block_size"))
+        def _parse_container_for_attr(container, attr):
+            s = set()
+            for a in container:
+                if type(a) is tuple or type(a) is list:
+                    s |= _parse_container_for_attr(a, attr)
+                else:
+                    if hasattr(a, attr):
+                        s.add(a.block_size)
+            return s
+
+        # Args is a tuple that can contain other tuples or lists. Parse it
+        # reccursively to find any BaseMCRTensor object
+        block_sizes = _parse_container_for_attr(args, "block_size")
+
         if len(block_sizes) != 1:
             raise RuntimeError(
                 f"Call to {func} must contain exactly one block size, got {list(block_sizes)}"