Candle equivalent to masked_scatter?

[vody-am]

Hi,

I'm currently porting some Pytorch code over to Candle, and ran into a stumbling block. The source makes use of masked_scatter, and there does not seem to be an equivalent in Candle:

https://github.com/huggingface/transformers/blob/main/src/transformers/models/paligemma/modeling_paligemma.py#L310-L311

  final_embedding = final_embedding.masked_scatter(
      image_mask.unsqueeze(-1).expand_as(final_embedding), scaled_image_features
  )

shape of final embedding: torch.Size([1, 260, 2048])

shape of image_mask torch.Size([1, 260])

shape of scaled features torch.Size([1, 256, 2048])

I would also be happy to route around this issue using a similar equivalent.

Example of masked_scatter:

self = torch.tensor([1, 2, 3, 4, 5, 6, 7])
mask = torch.tensor([1, 0, 0, 1, 0, 1, 1], dtype=torch.bool)
src = torch.tensor([100, 200, 300, 400])
self.masked_scatter(mask, src)
#  => tensor([100,   2,   3, 200,   5, 300, 400])

[vody-am]

cuda kernel looks to be this:

https://github.com/pytorch/pytorch/blob/99a11efc8ad8de16ffd538971a03398f6343c9a7/aten/src/ATen/native/cuda/IndexKernel.cu#L358-L409

CPU kernel looks to be this:

https://github.com/pytorch/pytorch/blob/99a11efc8ad8de16ffd538971a03398f6343c9a7/aten/src/ATen/native/cpu/IndexKernel.cpp#L336-L357

[LaurentMazare]

If the mask has a "regular" shape (e.g. the image embeddings are before or after the text ones in some dimensions), maybe you can just use a Tensor::cat instead?
Otherwise I would suggest processing this manually on the cpu to start with and get sure that the model works well.
Longer term, maybe add a custom op or maybe simpler: use index_select where the indexes are computed based on the mask itself + multiply by a 0/1 mask to ensure that only the relevant bits are kept.

[vody-am]

ah, cat might work. I'll dig into it. Thanks!

[vody-am]

just wanted to leave one more comment that it does more or less does boil down to a cat, with extra steps. Stepping back, this function (_merge_input_ids_with_image_features) roughly speaking does a concatenate of the 256 image "tokens" / embeddings, plus the text token embeddings into one big interleaved input in 3-D (with some logic for handling padding and the attention mask). Anyway I think I got past this!

[LaurentMazare]

Neat, very excited to see the model in action!