🐛[BUG]: ValueError: infer_schema(func) in @torch.library.custom_op due to stringified type annotations (from __future__ import annotations)

Version

e416350

On which installation method(s) does this occur?

Pip

Describe the issue

When importing physicsnemo (specifically through modules relying on physicsnemo.nn.functional.derivatives.mesh_green_gauss_gradient._warp_impl), a ValueError is raised during module load by PyTorch's schema inference mechanism in @torch.library.custom_op.

PyTorch's infer_schema inspects the function signature and expects actual types (such as <class 'torch.Tensor'>), but instead receives string literals (e.g. 'torch.Tensor'). This occurs when postponed evaluation of annotations is enabled (via from future import annotations or Python 3.11+ stringified annotations) without explicitly supplying the schema argument to @torch.library.custom_op.

Minimum reproducible example

python examples/structural_mechanics/deforming_plate/preprocessor.py

Relevant log output

D:\PycharmProjects\physicsnemo\venv\Scripts\python.exe D:\PycharmProjects\physicsnemo\examples\structural_mechanics\deforming_plate\preprocessor.py 
Warp CUDA warning: Could not find or load the NVIDIA CUDA driver. GPU execution will not be available.
Traceback (most recent call last):
  File "D:\PycharmProjects\physicsnemo\examples\structural_mechanics\deforming_plate\preprocessor.py", line 26, in <module>
    from deforming_plate_dataset import DeformingPlateDataset
  File "D:\PycharmProjects\physicsnemo\examples\structural_mechanics\deforming_plate\deforming_plate_dataset.py", line 31, in <module>
    from physicsnemo.datapipes.gnn.utils import load_json, save_json
  File "D:\PycharmProjects\physicsnemo\physicsnemo\datapipes\__init__.py", line 42, in <module>
    from physicsnemo.datapipes.dataset import Dataset
  File "D:\PycharmProjects\physicsnemo\physicsnemo\datapipes\dataset.py", line 40, in <module>
    from physicsnemo.datapipes.readers.base import Reader
  File "D:\PycharmProjects\physicsnemo\physicsnemo\datapipes\readers\__init__.py", line 29, in <module>
    from physicsnemo.datapipes.readers.mesh import DomainMeshReader, MeshReader
  File "D:\PycharmProjects\physicsnemo\physicsnemo\datapipes\readers\mesh.py", line 37, in <module>
    from physicsnemo.mesh import DomainMesh, Mesh
  File "D:\PycharmProjects\physicsnemo\physicsnemo\mesh\__init__.py", line 17, in <module>
    from physicsnemo.mesh.domain_mesh import DomainMesh
  File "D:\PycharmProjects\physicsnemo\physicsnemo\mesh\domain_mesh.py", line 28, in <module>
    from physicsnemo.mesh.mesh import Mesh, _requested_float_dtype
  File "D:\PycharmProjects\physicsnemo\physicsnemo\mesh\mesh.py", line 52, in <module>
    from physicsnemo.mesh.remeshing import remesh
  File "D:\PycharmProjects\physicsnemo\physicsnemo\mesh\remeshing\__init__.py", line 45, in <module>
    from physicsnemo.mesh.remeshing._partition import CellPartition, partition_cells
  File "D:\PycharmProjects\physicsnemo\physicsnemo\mesh\remeshing\_partition.py", line 42, in <module>
    from physicsnemo.nn.functional.neighbors import knn
  File "D:\PycharmProjects\physicsnemo\physicsnemo\nn\__init__.py", line 20, in <module>
    from .module.activations import (
  File "D:\PycharmProjects\physicsnemo\physicsnemo\nn\module\__init__.py", line 28, in <module>
    from .ball_query import BQWarp
  File "D:\PycharmProjects\physicsnemo\physicsnemo\nn\module\ball_query.py", line 29, in <module>
    from physicsnemo.nn.functional import radius_search
  File "D:\PycharmProjects\physicsnemo\physicsnemo\nn\functional\__init__.py", line 17, in <module>
    from .derivatives import (
  File "D:\PycharmProjects\physicsnemo\physicsnemo\nn\functional\derivatives\__init__.py", line 17, in <module>
    from .mesh_green_gauss_gradient import MeshGreenGaussGradient, mesh_green_gauss_gradient
  File "D:\PycharmProjects\physicsnemo\physicsnemo\nn\functional\derivatives\mesh_green_gauss_gradient\__init__.py", line 17, in <module>
    from .mesh_green_gauss_gradient import MeshGreenGaussGradient, mesh_green_gauss_gradient
  File "D:\PycharmProjects\physicsnemo\physicsnemo\nn\functional\derivatives\mesh_green_gauss_gradient\mesh_green_gauss_gradient.py", line 24, in <module>
    from ._warp_impl import mesh_green_gauss_gradient_warp
  File "D:\PycharmProjects\physicsnemo\physicsnemo\nn\functional\derivatives\mesh_green_gauss_gradient\_warp_impl.py", line 607, in <module>
    @torch.library.custom_op(
     ^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\PycharmProjects\physicsnemo\venv\Lib\site-packages\torch\_library\custom_ops.py", line 119, in inner
    schema_str = torch._custom_op.impl.infer_schema(fn, mutates_args)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\PycharmProjects\physicsnemo\venv\Lib\site-packages\torch\_library\infer_schema.py", line 42, in infer_schema
    error_fn(
  File "D:\PycharmProjects\physicsnemo\venv\Lib\site-packages\torch\_library\infer_schema.py", line 21, in error_fn
    raise ValueError(
ValueError: infer_schema(func): Parameter points has unsupported type torch.Tensor. The valid types are: dict_keys([<class 'torch.Tensor'>, typing.Optional[torch.Tensor], typing.Sequence[torch.Tensor], typing.List[torch.Tensor], typing.Sequence[typing.Optional[torch.Tensor]], typing.List[typing.Optional[torch.Tensor]], <class 'int'>, typing.Optional[int], typing.Sequence[int], typing.List[int], typing.Optional[typing.Sequence[int]], typing.Optional[typing.List[int]], <class 'float'>, typing.Optional[float], typing.Sequence[float], typing.List[float], typing.Optional[typing.Sequence[float]], typing.Optional[typing.List[float]], <class 'bool'>, typing.Optional[bool], typing.Sequence[bool], typing.List[bool], typing.Optional[typing.Sequence[bool]], typing.Optional[typing.List[bool]], <class 'str'>, typing.Optional[str], typing.Union[int, float, bool], typing.Union[int, float, bool, NoneType], typing.Sequence[typing.Union[int, float, bool]], typing.List[typing.Union[int, float, bool]], <class 'torch.dtype'>, typing.Optional[torch.dtype], <class 'torch.device'>, typing.Optional[torch.device]]). Got func with signature (points: 'torch.Tensor', cells: 'torch.Tensor', neighbors: 'torch.Tensor', values: 'torch.Tensor') -> 'torch.Tensor')

Process finished with exit code 1

Environment details