How to define custom ops with arbitrary number of tensors ?

[jeanfelixM]

trafficstars

I dived into the issues and repo about the subject, because I wanted to be able to call cuda kernels regarding 3D gaussian splatting, and the way to invoke those kernel seems to be custom ops. But right now, we only have

CustomOp1(Tensor, std::sync::Arc<Box<dyn CustomOp1 + Send + Sync>>),

CustomOp2(
        Tensor,
        Tensor,
        std::sync::Arc<Box<dyn CustomOp2 + Send + Sync>>,
    ),

CustomOp3(
        Tensor,
        Tensor,
        Tensor,
        std::sync::Arc<Box<dyn CustomOp3 + Send + Sync>>,
    )

And those gsplat kernels have way more in and/or out tensors depending on the operation.

I can think of ways to do it, but I was wondering if there was a good way to do it?

[jeanfelixM]

I'm trying to use the kernels of the gsplat project from nerfstudio, where for example, the kernel responsible of the forward pass of the gaussians projection is binded in C++ with libtorch with a function that has the signature :

std::tuple<
    torch::Tensor,
    torch::Tensor,
    torch::Tensor,
    torch::Tensor,
    torch::Tensor,
    torch::Tensor,
    torch::Tensor>
project_gaussians_forward_tensor(
    const int num_points,
    torch::Tensor &means3d,
    torch::Tensor &scales,
    const float glob_scale,
    torch::Tensor &quats,
    torch::Tensor &viewmat,
    torch::Tensor &projmat,
    const float fx,
    const float fy,
    const float cx,
    const float cy,
    const unsigned img_height,
    const unsigned img_width,
    const std::tuple<int, int, int> tile_bounds,
    const float clip_thresh
);

And i want to reproduce something like that in candle with that same kernel.

[jeanfelixM]

If anyone is interested, i used cat on input tensors of a forward method from a CustomOp, and narrow on the output Tensor. It kinda work, but I needed to change some function visibility.