Rewriting usage of torch.bucketize with more elementary functions

[EricLBuehler]

Hello everyone,

Thank you for your excellent work here. We are implementing the Idefics 2 model into mistral.rs which is built on Candle.

Candle has a Pytorch-like API, however, it does not have some functions. In particular, we noticed that Candle is missing the torch.bucketize function. I have opened an issue on Candle but thought it would also be a good idea to check here how to implement torch.bucketize using more elementary functions or generally rewriting that part of the code to emulate the behavior.

This is the part of the code I am referencing:

transformers/src/transformers/models/idefics2/modeling_idefics2.py

Lines 174 to 185 in 4b3eb19

	for batch_idx, p_attn_mask in enumerate(patch_attention_mask):
	nb_patches_h = p_attn_mask[:, 0].sum()
	nb_patches_w = p_attn_mask[0].sum()
	fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h)
	fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w)
	bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True)
	bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True)
	pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten()
	position_ids[batch_idx][p_attn_mask.view(-1).cpu()] = pos_ids

Is there a way to rewrite this such that I do not need to call torch.bucketize? Thank you for any help, I would really appreciate it!