Cannot compile torch.jit.script(kornia.geometry.warp_perspective) because of KORNIA_CHECK_IS_TENSOR

[josmople]

Describe the bug

The following should compile in torch script.
But there is an error because of type-checking.

import kornia
import torch


print("kornia.__version__", kornia.__version__)
print("torch.__version__", torch.__version__)

# kornia.core.check.KORNIA_CHECK_IS_TENSOR prevents compilation with torchscript
fn = torch.jit.script(kornia.geometry.warp_perspective)

Output and Error Message

kornia.__version__ 0.7.2
torch.__version__ 2.3.0
Traceback (most recent call last):
  File "/home/REDACTED/workspace/REDACTED/kornia_bug.py", line 9, in <module>
    fn = torch.jit.script(kornia.geometry.warp_perspective)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/REDACTED/anaconda3/envs/REDACTED/lib/python3.12/site-packages/torch/jit/_script.py", line 1395, in script
    fn = torch._C._jit_script_compile(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/REDACTED/anaconda3/envs/REDACTED/lib/python3.12/site-packages/torch/jit/_recursive.py", line 1007, in try_compile_fn
    return torch.jit.script(fn, _rcb=rcb)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/REDACTED/anaconda3/envs/REDACTED/lib/python3.12/site-packages/torch/jit/_script.py", line 1395, in script
    fn = torch._C._jit_script_compile(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/REDACTED/anaconda3/envs/REDACTED/lib/python3.12/site-packages/torch/jit/_recursive.py", line 1007, in try_compile_fn
    return torch.jit.script(fn, _rcb=rcb)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/REDACTED/anaconda3/envs/REDACTED/lib/python3.12/site-packages/torch/jit/_script.py", line 1395, in script
    fn = torch._C._jit_script_compile(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: 
Unknown type name 'object':
  File "/home/REDACTED/anaconda3/envs/REDACTED/lib/python3.12/site-packages/kornia/core/check.py", line 149
def KORNIA_CHECK_IS_TENSOR(x: object, msg: Optional[str] = None, raises: bool = True) -> TypeGuard[Tensor]:
                              ~~~~~~ <--- HERE
    """Check the input variable is a Tensor.
'KORNIA_CHECK_IS_TENSOR' is being compiled since it was called from 'transform_points'
  File "/home/REDACTED/anaconda3/envs/REDACTED/lib/python3.12/site-packages/kornia/geometry/linalg.py", line 174
        >>> points_0 = transform_points(trans_01, points_1)  # BxNx3
    """
    KORNIA_CHECK_IS_TENSOR(trans_01)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    KORNIA_CHECK_IS_TENSOR(points_1)
    if not trans_01.shape[0] == points_1.shape[0] and trans_01.shape[0] != 1:
'transform_points' is being compiled since it was called from 'warp_perspective'
  File "/home/REDACTED/anaconda3/envs/REDACTED/lib/python3.12/site-packages/kornia/geometry/transform/imgwarp.py", line 124
        .expand(B, h_out, w_out, 2)
    )
    grid = transform_points(src_norm_trans_dst_norm[:, None, None], grid)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

    if padding_mode == "fill":

I tried commenting out KORNIA_CHECK_IS_TENSOR in kornia.geometry.linalg

def transform_points(trans_01: Tensor, points_1: Tensor) -> Tensor:
    # KORNIA_CHECK_IS_TENSOR(trans_01)
    # KORNIA_CHECK_IS_TENSOR(points_1)
   ...

And torch script compiles.

Reproduction steps

1. import kornia
2. import torch
3. torch.jit.script(kornia.geometry.warp_perspective)

Expected behavior

It should compile without any errors.

Environment

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] onnx==1.16.0
[pip3] pytorch-lightning==2.2.1
[pip3] torch==2.3.0
[pip3] torchaudio==2.3.0
[pip3] torchmetrics==1.3.2
[pip3] torchvision==0.18.0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] mkl-service               2.4.0           py312h5eee18b_1  
[conda] mkl_fft                   1.3.8           py312h5eee18b_0  
[conda] mkl_random                1.2.4           py312hdb19cb5_0  
[conda] numpy                     1.26.4          py312hc5e2394_0  
[conda] numpy-base                1.26.4          py312h0da6c21_0  
[conda] pytorch                   2.3.0           py3.12_cuda12.1_cudnn8.9.2_0    pytorch
[conda] pytorch-cuda              12.1                 ha16c6d3_5    pytorch
[conda] pytorch-lightning         2.2.1                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch                     2.3.0                    pypi_0    pypi
[conda] torchaudio                2.3.0               py312_cu121    pytorch
[conda] torchmetrics              1.3.2                    pypi_0    pypi
[conda] torchvision               0.18.0              py312_cu121    pytorch

Additional context

No response

[johnnv1]

related to #2559

thanks @josmople for opening this issue, the kornia-core API sadly isn't working within JIT at the moment, and we expect to fix it in the future, open to contributions too

[edgarriba]

maybe we could include torch.jit.is_srcipting / torch.jit.is_tracing in the checker api
https://pytorch.org/docs/master/jit_language_reference.html#torch.jit.is_scripting