Not sure if this is comparable to dashesy’s in terms of generality, but here’s a simplified case + tests for selecting vectors, e.g. via argmin or multinomial:

  

      gist.github.com
  

  
    https://gist.github.com/EricCousineau-TRI/cc2dc27c7413ea8e5b4fd9675050b1c0 22

  vector_gather.py
  import einops
import torch

def vector_gather(vectors, indices):
    """
    Gathers (batched) vectors according to indices.

    Arguments:
        vectors: Tensor[N, L, D]
        indices: Tensor[N, K] or Tensor[N]
    This file has been truncated. show original
  vector_gather_test.py
  import unittest

import einops
import torch


class Test(unittest.TestCase):
    def test_vector_gather(self):
        vector = torch.tensor([
            [
    This file has been truncated. show original