Match Behavior of mn.transforms.Affined and mn.networks.layers.AffineTransform

[sudomakeinstall]

First, thanks to everyone for maintaining this incredible project. :-)

I'm trying to understand a difference in behavior between the transform mn.transforms.Affined and the layer mn.networks.layers.AffineTransform; specifically, if I apply a transform A using mn.transforms.Affined in my dataloader and then subsequently try to reverse this transformation within my network using mn.networks.layers.AffineTransform, the images do not align. I suspect that these classes are making different assumptions about normalization and/or orientation, but am having trouble parsing this out from the documentation [1,2].

Below is a minimal example reproducing the issue--could anyone suggest how I could modify this so that the output aligns with the input (after resampling/cropping but before applying the affine transform)? Note that in my actual application I can't use the .invert() function because I'm trying to learn the transformation parameters that I'm applying.

Best, and thanks,

--Davis

[1] https://docs.monai.io/en/stable/transforms.html#monai.transforms.Affined
[2] https://docs.monai.io/en/stable/networks.html#affine-transform

#!/usr/bin/env python

import monai as mn
import torch as t

data_dict = {"img": "/path/to/nifti.nii.gz"}

A = t.tensor([
    [ 0.6253, -0.4745, -0.6196,  0.0000],
    [-0.5303, -0.8408,  0.1087,  0.0000],
    [-0.5725,  0.2606, -0.7774,  0.0000],
    [ 0.0000,  0.0000,  0.0000,  1.0000],
    ])

load = mn.transforms.LoadImaged(
    keys=("img"),
    image_only=True,
    ensure_channel_first=True,
    simple_keys=True,
)
orient = mn.transforms.Orientationd(
    keys=("img"),
    axcodes="LAS",
)
space = mn.transforms.Spacingd(
    keys=("img"),
    pixdim=[1.0 for _ in range(3)],
    mode="bilinear",
)
crop = mn.transforms.ResizeWithPadOrCropd(
    keys=("img"),
    spatial_size=(128, 128, 128),
)
affine_fw = mn.transforms.Affined(
    keys=("img"),
    affine=A,
    mode="bilinear",
    padding_mode="zeros",
)
unsqueeze = mn.transforms.Lambdad(
    keys=("img"),
    func=lambda x: x.unsqueeze(0),
)
affine_rv = mn.networks.layers.AffineTransform(
        zero_centered=True,
)

pipeline_1 = mn.transforms.Compose(
    [
        load,
        orient,
        space,
        crop,
        affine_fw,
        unsqueeze,
    ]
)

data_dict = pipeline_1(data_dict)
data_dict["img"] = affine_rv(data_dict["img"], t.linalg.inv(A))