tensor: torch.Tensor of shape [batch_size, seq_len] C: scalar value to match Returns: torch.Tensor of shape [batch_size] with last indices
(tensor, C)
| 146 | |
| 147 | |
| 148 | def find_last_equal_C(tensor, C): |
| 149 | """ |
| 150 | tensor: torch.Tensor of shape [batch_size, seq_len] |
| 151 | C: scalar value to match |
| 152 | Returns: torch.Tensor of shape [batch_size] with last indices |
| 153 | """ |
| 154 | mask = (tensor == C).int() # Shape: [batch_size, seq_len], bool tensor |
| 155 | flipped_mask = mask.flip(dims=[1]) # Flip along sequence dimension |
| 156 | flipped_indices = flipped_mask.argmax(dim=1) # First True in flipped |
| 157 | seq_len = tensor.shape[1] |
| 158 | last_indices = (seq_len - 1) - flipped_indices # Convert to original indices |
| 159 | |
| 160 | # Optional: Handle cases with no C (set to -1), though problem assumes existence |
| 161 | actual_values = tensor[torch.arange(tensor.shape[0]), last_indices] |
| 162 | no_match = actual_values != C |
| 163 | last_indices[no_match] = -1 |
| 164 | |
| 165 | return last_indices |