Additional Tensor Operations Needed

[MilkFather]

I am re-implementing the famous NeRF (https://github.com/bmild/nerf) model in candle. During my experience with candle, I found the following Tensor operations missing.

• torch.linspace
• torch.cumprod
• torch.searchsorted
• torch.sort

I am able to find a workaround for torch.linspace and torch.cumprod, but torch.searchsorted and torch.sort look formidable. Are there recommended ways to implement that efficiently? Thanks.

Current workarounds for linspace and cumprod

• linspace: directly taken from the stable-diffusion example.
• cumprod: modified from https://github.com/mokeyish/candle-ext/blob/1b38cb1b9faa9694c204d46788fd70d479109bab/src/cumsum.rs#L6-L28 changing the add operation to multiply operation. Not the optimal solution, I admit.

[EricLBuehler]

Perhaps see #1351? I also noticed there are definitely some missing Tensor features.

[LaurentMazare]

For cumprod, if all your values are positive I would recommend doing log then the built-in cumsum then exp, this should be more efficient and easier than going the row data.
For searchsorted and sort, I think you should do a custom-op for these. The goal of custom ops is to have candle be extensible so that operations that are not super common can be added by the user. You can see an example of implementation here.
Let me know if you have further questions or you want me to sketch a basic sort function using this.

[MilkFather]

I would like to see an example implementation of sort.

The sort function in PyTorch returns a tuple of tensors - one for the sorted tensor (preserving the original DType) and another one for indexes (a LongTensor that can be used for indexing).

From what I understand, the custom op traits have a fixed output type (Result<(Storage, Shape)>) that looks incompatible with this particular use case. Hence I would like to see an example implementation for that.

[LaurentMazare]

Or a simple idea there would just be to have a sort_index custom op that return the sorted indexes and to use some gather function to get the sorted values if needed.

[MilkFather]

@LaurentMazare I spent the last weekend digging through the documentation and found that the crate is very under-documented. The example given does not touch layout manipulation at all and does little help to understanding the required steps to implement a custom operation like sort or search.

By the way, none of the "some gathering function" really pops into my mind, which leads me back to the starting point.

EDIT: I just noticed that Tensor::gather() does exactly what I needed. I can write an extension trait for Tensor.

impl TensorExt for Tensor {
    fn sort<D: dim>(&self, dim: D, ascending: bool) -> Result<(Tensor, Tensor)> {
        let sorted_idx = self.apply_op1(/* something */)?;
        let sorted_val = self.gather(&sorted_idx, dim)?;
        Ok((sorted_val, sorted_idx))
    }
}

Although I still have no idea how to implement that /* something */.

[LaurentMazare]

I've tweaked the basics example to show how to implement the custom op. It's only implemented for contiguous tensors (so no need to care about the layout) and only for f32 but should be easy to extend if you want to cover other dtypes.
Code here. And yes documentation is lacking, we've put a lot of effort working on this crate but have focused on applications so far and a lot of things have to be polished.

[jeromeku]

@LaurentMazare @MilkFather

Added a candle-example of a CustomOp2 that implements torch.searchsorted.

See PR #1418.

[MilkFather]

Sorry for not responding for a long time, I am too occupied in real life. It is a thrill to see your (and the community's) great work. I am grateful to it.