Fix(argo2): Restore compatibility with modern kornia versions (>=0.7.0)

[Qu0rise]

Subject: Ensuring Argoverse 2 usability with current kornia releases

Dear Maintainers,

This pull request resolves a compatibility issue preventing the Argoverse 2 dataset utilities from functioning correctly with kornia version 0.7.0 and later.

Motivation & Impact:

The core issue stems from kornia removing a deprecated order argument in its rotation conversion API (v0.7.0+). Our current codebase in pcdet/datasets/argo2/argo2_utils/so3.py still utilizes this argument, leading to runtime errors for users who maintain up-to-date dependencies. This negatively impacts the user experience for those working with Argoverse 2 and potentially increases support requests related to environment setup (similar context to closed issue #1470).

Merging this PR offers several key benefits:

• Restores Functionality: Enables seamless use of Argoverse 2 features for users with modern kornia installations.
• Enhances Robustness: Eliminates a known source of errors, making the codebase more resilient to dependency updates.
• Reduces Maintenance: Proactively addresses a compatibility issue, potentially preventing future bug reports.

Technical Solution:

The fix is straightforward: it removes the now-obsolete order=C.QuaternionCoeffOrder.WXYZ argument from the quaternion_to_rotation_matrix and rotation_matrix_to_quaternion function calls within pcdet/datasets/argo2/argo2_utils/so3.py. This aligns our code with the current kornia API without altering the intended WXYZ (scalar-first) quaternion convention, which remains the default or expected behavior in recent kornia versions.

Review Focus & Testing:

This change is highly localized to the argo2 utilities and directly addresses an external library API change.

• The core logic of the rotation conversions remains unchanged.
• As kornia is an optional dependency primarily for specific datasets like argo2, the risk to other parts of OpenPCDet is minimal.
• Manual testing confirms the fix resolves the error with kornia>=0.7.0. Due to the nature of the fix (API compatibility), extensive new unit tests were deemed unnecessary for this specific change.

This contribution aims to improve the project's usability and maintainability with minimal disruption. I appreciate your time reviewing this change and welcome any feedback.

Best regards,

@Qu0rise

Related to: #1470

[liyufan]

According to kornia/kornia#2437, the (x, y, z, w) order was deprecated and the new default order is (w, x, y, z), so it's safe to remove the order argument here. Moreover, please consider adding kornia and other missing dependencies in requirements.txt:

av2
kornia>=0.7
ninja

[liyufan]

It seems that the @torch.jit.script above function quat_to_mat should be removed otherwise it will raise exception.

[liyufan]

When generate the data infos:

python -m pcdet.datasets.kitti.kitti_dataset create_kitti_infos tools/cfgs/dataset_configs/kitti_dataset.yaml

The output is:

Traceback (most recent call last):
  File "<frozen runpy>", line 189, in _run_module_as_main
  File "<frozen runpy>", line 112, in _get_module_details
  File "/home/lyf/openpcdet/pcdet/datasets/__init__.py", line 15, in <module>
    from .argo2.argo2_dataset import Argo2Dataset
  File "/home/lyf/openpcdet/pcdet/datasets/argo2/argo2_dataset.py", line 15, in <module>
    from .argo2_utils.so3 import yaw_to_quat, quat_to_yaw
  File "/home/lyf/openpcdet/pcdet/datasets/argo2/argo2_utils/so3.py", line 9, in <module>
    @torch.jit.script
     ^^^^^^^^^^^^^^^^
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/torch/jit/_script.py", line 1443, in script
    ret = _script_impl(
          ^^^^^^^^^^^^^
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/torch/jit/_script.py", line 1214, in _script_impl
    fn = torch._C._jit_script_compile(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/torch/jit/_recursive.py", line 993, in try_compile_fn
    return torch.jit.script(fn, _rcb=rcb)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/torch/jit/_script.py", line 1443, in script
    ret = _script_impl(
          ^^^^^^^^^^^^^
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/torch/jit/_script.py", line 1214, in _script_impl
    fn = torch._C._jit_script_compile(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: 
cannot statically infer the expected size of a list in this context:
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/kornia/geometry/conversions.py", line 586

    # this slightly awkward construction of the output shape is to satisfy torchscript
    output_shape = [*list(quaternion.shape[:-1]), 3, 3]
                    ~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    matrix = matrix_flat.reshape(output_shape)
'quaternion_to_rotation_matrix' is being compiled since it was called from 'quat_to_mat'
  File "/home/lyf/openpcdet/pcdet/datasets/argo2/argo2_utils/so3.py", line 19
        (...,3,3) 3D rotation matrices.
    """
    return C.quaternion_to_rotation_matrix(quat_wxyz)
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

When run the demo:

python demo.py --cfg_file cfgs/custom_models/PartA2_free.yaml  --ckpt ../output/custom_models/PartA2_free/default/ckpt/checkpoint_epoch_500.pth  --data_path ../data/custom/points/0020.npy --ext .npy

The output is:

Traceback (most recent call last):
  File "/home/lyf/openpcdet/tools/demo.py", line 18, in <module>
    from pcdet.datasets import DatasetTemplate
  File "/home/lyf/openpcdet/pcdet/datasets/__init__.py", line 15, in <module>
    from .argo2.argo2_dataset import Argo2Dataset
  File "/home/lyf/openpcdet/pcdet/datasets/argo2/argo2_dataset.py", line 15, in <module>
    from .argo2_utils.so3 import yaw_to_quat, quat_to_yaw
  File "/home/lyf/openpcdet/pcdet/datasets/argo2/argo2_utils/so3.py", line 9, in <module>
    @torch.jit.script
     ^^^^^^^^^^^^^^^^
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/torch/jit/_script.py", line 1443, in script
    ret = _script_impl(
          ^^^^^^^^^^^^^
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/torch/jit/_script.py", line 1214, in _script_impl
    fn = torch._C._jit_script_compile(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/torch/jit/_recursive.py", line 993, in try_compile_fn
    return torch.jit.script(fn, _rcb=rcb)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/torch/jit/_script.py", line 1443, in script
    ret = _script_impl(
          ^^^^^^^^^^^^^
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/torch/jit/_script.py", line 1214, in _script_impl
    fn = torch._C._jit_script_compile(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: 
cannot statically infer the expected size of a list in this context:
  File "/home/lyf/anaconda3/envs/py/lib/python3.12/site-packages/kornia/geometry/conversions.py", line 586

    # this slightly awkward construction of the output shape is to satisfy torchscript
    output_shape = [*list(quaternion.shape[:-1]), 3, 3]
                    ~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    matrix = matrix_flat.reshape(output_shape)
'quaternion_to_rotation_matrix' is being compiled since it was called from 'quat_to_mat'
  File "/home/lyf/openpcdet/pcdet/datasets/argo2/argo2_utils/so3.py", line 19
        (...,3,3) 3D rotation matrices.
    """
    return C.quaternion_to_rotation_matrix(quat_wxyz)
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

So I assume that the decorator @torch.jit.script of function quat_to_mat should be removed. Note that the decorator @torch.jit.script of the following function mat_to_quat is commented out, so it is reasonable to remove it too.

[liyufan]

The decorator @torch.jit.script of this function should be removed.

[liyufan]

Actually, the error is caused by kornia dropping JIT support in favor of TorchDynamo. Related links:

• Dropping JIT support kornia/kornia#2200
• drop JIT support from geometry.conversions kornia/kornia#2424

The code change that causes the error is: kornia/kornia@e98fd8c#diff-a8f3e0de21b707884a7a9a8a02a24848470a376b1b49bc585f7a88d8794f167bR584. While the original code could work:

output_shape = list(quaternion.shape[:-1]) + [3, 3]
# or
output_shape = quaternion.shape[:-1] + (3, 3)

The commit changed it to:

output_shape = [*list(quaternion.shape[:-1]), 3, 3]

which broke static shape inference of TorchScript due to the dynamic list expansion used to construct output_shape. So the decorator @torch.jit.script must be removed here.

[Qu0rise]

Thanks for the detailed investigation, @liyufan!

What I changed

• Removed @torch.jit.script from the two kornia wrappers

  • quat_to_mat — decorator dropped
  • mat_to_quat — kept plain Python (wrapper was already unscripted)
    Kornia ≥ 0.7 abandons TorchScript in favour of TorchDynamo, so scripting
    these functions triggers the shape-inference error you reported.

• Left the remaining four helpers scripted
  They are pure PyTorch math, compile cleanly under TorchScript and do not
  touch kornia. If you prefer full consistency I can drop those decorators as
  well—just let me know.