text: #!/usr/bin/env python3
  
  import torch
  
  from captum.attr._utils.batching import (
      _batched_generator,
      _batched_operator,
      _tuple_splice_range,
  )
  from tests.helpers.basic import BaseTest, assertTensorAlmostEqual
  
  
  class Test(BaseTest):
      def test_tuple_splice_range(self):
          test_tuple = (
              torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]]),
              "test",
              torch.tensor([[6, 7, 8], [0, 1, 2], [3, 4, 5]]),
          )
          spliced_tuple = _tuple_splice_range(test_tuple, 1, 3)
          assertTensorAlmostEqual(self, spliced_tuple[0], [[3, 4, 5], [6, 7, 8]])
          self.assertEqual(spliced_tuple[1], "test")
          assertTensorAlmostEqual(self, spliced_tuple[2], [[0, 1, 2], [3, 4, 5]])
  
      def test_tuple_splice_range_3d(self):
          test_tuple = (
              torch.tensor([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [6, 7, 8]]]),
              "test",
          )
          spliced_tuple = _tuple_splice_range(test_tuple, 1, 2)
          assertTensorAlmostEqual(self, spliced_tuple[0], [[[6, 7, 8], [6, 7, 8]]])
          self.assertEqual(spliced_tuple[1], "test")
  
      def test_batched_generator(self):
          def sample_operator(inputs, additional_forward_args, target_ind, scale):
              return (
                  scale * (sum(inputs)),
                  scale * sum(additional_forward_args),
                  target_ind,
              )
  
          array1 = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
          array2 = [[6, 7, 8], [0, 1, 2], [3, 4, 5]]
          array3 = [[0, 1, 2], [0, 0, 0], [0, 0, 0]]
          inp1, inp2, inp3 = (
              torch.tensor(array1),
              torch.tensor(array2),
              torch.tensor(array3),
          )
          for index, (inp, add, targ) in enumerate(
              _batched_generator((inp1, inp2), (inp3, 5), 7, 1)
          ):
              assertTensorAlmostEqual(self, inp[0], array1[index])
              assertTensorAlmostEqual(self, inp[1], array2[index])
              assertTensorAlmostEqual(self, add[0], array3[index])
              self.assertEqual(add[1], 5)
              self.assertEqual(targ, 7)
  
      def test_batched_operator_0_bsz(self):
          inp1 = torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
          with self.assertRaises(AssertionError):
              _batched_operator(lambda x: x, inputs=inp1, internal_batch_size=0)
  
      def test_batched_operator(self):
          def _sample_operator(inputs, additional_forward_args, target_ind, scale):
              return (
                  scale * (sum(inputs)),
                  scale * sum(additional_forward_args) + target_ind[0],
              )
  
          inp1 = torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
          inp2 = torch.tensor([[6, 7, 8], [0, 1, 2], [3, 4, 5]])
          inp3 = torch.tensor([[0, 1, 2], [0, 0, 0], [0, 0, 0]])
          batched_result = _batched_operator(
              _sample_operator,
              inputs=(inp1, inp2),
              additional_forward_args=(inp3),
              target_ind=[0, 1, 2],
              scale=2.0,
              internal_batch_size=1,
          )
          assertTensorAlmostEqual(
              self, batched_result[0], [[12, 16, 20], [6, 10, 14], [18, 22, 26]]
          )
          assertTensorAlmostEqual(
              self, batched_result[1], [[0, 2, 4], [1, 1, 1], [2, 2, 2]]
          )
  