basis.py 389 Bytes
Newer Older
rusty1s's avatar
rusty1s committed
1
from typing import Tuple
rusty1s's avatar
rusty1s committed
2

rusty1s's avatar
rusty1s committed
3
import torch
rusty1s's avatar
rusty1s committed
4
5


rusty1s's avatar
rusty1s committed
6
7
8
9
10
11
@torch.jit.script
def spline_basis(pseudo: torch.Tensor, kernel_size: torch.Tensor,
                 is_open_spline: torch.Tensor,
                 degree: int) -> Tuple[torch.Tensor, torch.Tensor]:
    return torch.ops.torch_spline_conv.spline_basis(pseudo, kernel_size,
                                                    is_open_spline, degree)