Usage of | in type annotation causes torch.jit.script to not work

[balbok0]

Describe the bug

torch.jit.script on a model that calls warp_affine within forward function fails with following error:

RuntimeError: 
Expression of type | cannot be used in a type expression:
  File "/usr/local/lib/python3.8/dist-packages/kornia/geometry/conversions.py", line 1015
def normal_transform_pixel(
    height: int, width: int, eps: float = 1e-14, device: torch.device | None = None, dtype: torch.dtype | None = None
                                                         ~~~~~~~~~~~~~~~~~~~ <--- HERE
) -> Tensor:
    r"""Compute the normalization matrix from image size in pixels to [-1, 1].
'normal_transform_pixel' is being compiled since it was called from 'normalize_homography'
  File "/usr/local/lib/python3.8/dist-packages/kornia/geometry/conversions.py", line 1004

    # compute the transformation pixel/norm for src/dst
    src_norm_trans_src_pix: Tensor = normal_transform_pixel(src_h, src_w).to(dst_pix_trans_src_pix)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

    src_pix_trans_src_norm = _torch_inverse_cast(src_norm_trans_src_pix)
'normalize_homography' is being compiled since it was called from 'warp_affine'
  File "/usr/local/lib/python3.8/dist-packages/kornia/geometry/transform/imgwarp.py", line 201
    # we generate a 3x3 transformation matrix from 2x3 affine
    M_3x3: Tensor = convert_affinematrix_to_homography(M)
    dst_norm_trans_src_norm: Tensor = normalize_homography(M_3x3, (H, W), dsize)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

    # src_norm_trans_dst_norm = torch.inverse(dst_norm_trans_src_norm)
'warp_affine' is being compiled since it was called from 'DummyModel.forward'
  File "tmp.py", line 6
    def forward(self, x, y):
        return T.warp_affine(
               ~~~~~~~~~~~~~~
            x, y,
            ~~~~~
            dsize=(1920, 1080),
            ~~~~~~~~~~~~~~~~~~~
            align_corners=False,
            ~~~~~~~~~~~~~~~~~~~ <--- HERE
        )

Reproduction steps

import torch
from kornia.geometry import transform as T

class DummyModel(torch.nn.Module):
    def forward(self, x, y):
        return T.warp_affine(
            x, y,
            dsize=(1920, 1080),
            align_corners=False,
        )

torch.jit.script(DummyModel())

Expected behavior

torch.jit.script compiles.

Environment

wget https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py

• PyTorch Version (e.g., 1.0): 2.0.1+cu117 (also 2.0.1)
• OS (e.g., Linux): Linux
• How you installed PyTorch (conda, pip, source): pip
• Build command you used (if compiling from source):
• Python version: 3.8 (also happening in 3.9)
• CUDA/cuDNN version: 11.7
• GPU models and configuration: Tested on 3080, 4090 and CPU
• Any other relevant information:

Additional context

Optional[torch.device] will work with torch.jit.script. I understand that torch.device | None is cleaner code, and more pythonic, but it would be great if this library supported jit

[johnnv1]

This is expected because we started to drop jit support (see #2200) in favor of torch compile -- but we'll gradually get back to it see #2559. Could you submit a PR to fix it on the geometries file?

[edgarriba]

ideally we should verify that this can be exported to onnx as well

[balbok0]

I can try doing that. Thanks for pointing me to torch compile though, I'm still catching up on torch>=2.0 features!

For verifying ONNX support, are there any specific functions/classes that I should verify against? I'm not super familiar with kornia lib as a whole.

[johnnv1]

You can find any test_onnx we had in the test suite, and check updating it with bout the annotations styles