code: def test_index_select_cat(D, batches) -> None:
      torch.manual_seed(0)
      dtype = torch.float16
  
      num_rows = 0
      for B, seqlen in batches:
          num_rows += B * seqlen
  
      src = torch.randn(
          [num_rows, D], device="cuda", dtype=dtype, requires_grad=True
      )
      indices = []
      sources = []
      rows_begin = 0
      for B, seqlen in batches:
          index = [i for i in range(B)]
          random.Random(B).shuffle(index)
          indices.append(
              torch.tensor(
                  index[: int(0.6 * B)],
                  dtype=torch.int64,
                  device="cuda",
              )
          )
          sources.append(
              src[rows_begin : rows_begin + B * seqlen].reshape(
                  [B, seqlen * D]
              )
          )
          rows_begin += B * seqlen
  
      # PT implem
      ref_out = torch.cat(
          [s[i].flatten() for s, i in zip(sources, indices)], dim=0
      )
      gradient_out = torch.randn_like(ref_out)
      ref_out.backward(gradient_out)
      assert src.grad is not None
      ref_grad = src.grad.clone()
      src.grad = None
  
      # xFormers implem
      out = xops.index_select_cat(sources, indices)
      assert_allclose(out, ref_out, "fw")
      out.backward(gradient_out)
      assert src.grad is not None
      assert_allclose(src.grad, ref_grad, "src.grad")
  
file_path: xformers/tests/test_indexing.py
imports: import random
  import pytest
  import torch
  import xformers.ops
  from xformers.ops import indexing
  from .utils import assert_allclose
label_desc: Use image node_id 2 for calling a global function with example usage: test_index_select_cat(D, batches) without return types