metatensor einops rearrange

[wyli]

Describe the bug

import torch
from monai.data import MetaTensor
from einops import rearrange
from monai.transforms import Resize

# create a metatensor in batch mode
x = MetaTensor(torch.randn((1, 1, 481, 404, 3))).to("cuda:0")
x.is_batch = True
x.meta['affine'] = torch.eye(4)[None]

# select the first element in the batch and resize
out = Resize(spatial_size=(1024, 860, -1), align_corners=True, mode="trilinear")(x[0])
print(out.shape)  # works fine

# rearrange using einops
x_ = rearrange(x, "b c h w d -> (b c) h w d")
out = Resize(spatial_size=(1024, 860, -1), align_corners=True, mode="trilinear")(x_)
print(out.shape)  # error

error:

    affine = orig_affine @ to_affine_nd(len(orig_affine) - 1, affine, dtype=torch.float64)
ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 1 is different from 4)

workaround (convert to a regular tensor as_tensor()):

x_ = rearrange(x, "b c h w d -> (b c) h w d")
if isinstance(x_, MetaTensor):
    x_ = x_.as_tensor()
out = Resize(spatial_size=(1024, 860, -1), align_corners=True, mode="trilinear")(x_)
print(out.shape)

[aymuos15]

I want to pick up this if still relevant. Could someone give a starting guide as to how I could approach this?