Fix dtype propagation

[janEbert]

Previously, dtypes were not propagated correctly in composed transforms, resulting in errors when different dtypes were mixed. Even specifying a dtype in the constructor does not fix this. Neither does specifying the dtype for each composition function invocation (e.g. as a kwarg in rotate_axis_angle).

With the change, I also had to modify the default dtype of RotateAxisAngle, which was torch.float64; it is now torch.float32 like for all other transforms. This was required because the fix in propagation broke some tests due to dtype mismatches.

This change in default dtype in turn broke two tests due to precision changes (calculations that were previously done in torch.float64 were now done in torch.float32), so I changed the precision tolerances to be less strict. I chose the lowest power of ten that passed the tests here.

[janEbert]

For someone that stumbles upon the bug and needs an immediate solution, here's a monkey patch:

import pytorch3d.transforms as tfs

def fix_pytorch3d() -> None:
    """Monkey patch missing dtype propagation using default dtype."""
    # Neither a good nor elegant solution, but it handles the bug.
    tfs.Transform3d.__init__.__defaults__ = (
        (th.get_default_dtype(),)
        + tfs.Transform3d.__init__.__defaults__[1:]
    )
    tfs.Translate.__init__.__defaults__ = (
        tfs.Translate.__init__.__defaults__[:2]
        + (th.get_default_dtype(),)
        + tfs.Translate.__init__.__defaults__[3:]
    )
    tfs.Scale.__init__.__defaults__ = (
        tfs.Scale.__init__.__defaults__[:2]
        + (th.get_default_dtype(),)
        + tfs.Scale.__init__.__defaults__[3:]
    )
    tfs.Rotate.__init__.__defaults__ = (
        (th.get_default_dtype(),)
        + tfs.Rotate.__init__.__defaults__[1:]
    )
    tfs.RotateAxisAngle.__init__.__defaults__ = (
        tfs.Scale.__init__.__defaults__[:2]
        + (th.get_default_dtype(),)
        + tfs.Rotate.__init__.__defaults__[3:]
    )

[facebook-github-bot]

@bottler has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[bottler]

Thank you for fixing this!