code: def batch_add(tensor, idx, val, device=None):
      """Do addition per row (first axis)."""
      assert tensor.shape[0] == idx.shape[0]  # Same size of first dim
      dim0_size, dim1_size = tensor.shape[0], tensor.shape[1]
  
      if device is None:
          device = tensor.device
  
      tensor = torch.reshape(tensor, [dim0_size * dim1_size, -1])
      idx_offset = torch.arange(0, dim0_size, device=device) * dim1_size
      for _ in range(len(idx.size()) - 1):
          idx_offset = idx_offset.unsqueeze(-1)
      new_idx = idx + idx_offset
  
      val = val.reshape(val.size(0) * val.size(1), -1)
      res = tensor.index_add(0, new_idx.view(-1), val).reshape(
          [dim0_size, dim1_size, -1]
      )
  
      if (
          tensor.shape[-1] == 1
      ):  # If selected element is scalar, restore original dim
          res = res.squeeze(-1)
  
      return res
  
imports: import sys
  import os
  from os import makedirs
  from os.path import join
  import numpy
  import pyhocon
  import math
  import logging
  import random
  from typing import Optional, Tuple, Any, Dict, Iterable, Union
  import torch
  import torch.nn.functional
  from torch import nn
  from transformers.activations import ACT2FN
label_desc: Use image node_id 7 for calling a global function with example usage: batch_add(tensor, idx, val, device) and returns: res