for (int batch_idx : llvm::seq<int>(0, bcast.output_batch_size())) {
    int lhs_batch_idx, rhs_batch_idx;
    if (bcast.IsBroadcastingRequired()) {
      lhs_batch_idx = bcast.x_batch_indices()[batch_idx];
      rhs_batch_idx = bcast.y_batch_indices()[batch_idx];
    } else {
      lhs_batch_idx = batch_idx;
      rhs_batch_idx = batch_idx;
    }
    auto matmul = rewriter.create<TF::MatMulOp>(loc, matmul_type,

         Output(Tile(scope, one,
                     ExpandDims(scope, Sub(scope, indices_ndims, dim), 0)))},
        /*axis=*/0);
    batch_indices =
        Add(scope, batch_indices, Reshape(scope, dim_indices, dim_shape));
  }

  return batch_indices;
}

Output BatchGatherGrad(const Scope& scope, Output params_shape, Output values,