class_docstr: This class tests all functions of the ArithmeticSharedTensor.
docstr: Tests arithmetic functions on encrypted tensor.
docstr_tok: ["Tests","arithmetic","functions","on","encrypted","tensor","."]
code: def test_arithmetic(self) -> None:
      """Tests arithmetic functions on encrypted tensor."""
      arithmetic_functions = [
          "add",
          "add_",
          "sub",
          "sub_",
          "mul",
          "mul_",
      ]
      for func in arithmetic_functions:
          for tensor_type in [lambda x: x, ArithmeticSharedTensor]:
              tensor1 = get_random_test_tensor(is_float=True)
              tensor2 = get_random_test_tensor(is_float=True)
              encrypted = ArithmeticSharedTensor(tensor1)
              encrypted2 = tensor_type(tensor2)
  
              reference = getattr(tensor1, func)(tensor2)
              encrypted_out = getattr(encrypted, func)(encrypted2)
              private_type = tensor_type == ArithmeticSharedTensor
              self._check(
                  encrypted_out,
                  reference,
                  "%s %s failed"
                  % ("private" if private_type else "public", func),
              )
              if "_" in func:
                  # Check in-place op worked
                  self._check(
                      encrypted,
                      reference,
                      "%s %s failed"
                      % ("private" if private_type else "public", func),
                  )
              else:
                  # Check original is not modified
                  self._check(
                      encrypted,
                      tensor1,
                      "%s %s failed"
                      % (
                          "private"
                          if tensor_type == ArithmeticSharedTensor
                          else "public",
                          func,
                      ),
                  )
  
              # Check encrypted vector with encrypted scalar works.
              tensor1 = get_random_test_tensor(is_float=True)
              tensor2 = get_random_test_tensor(is_float=True, size=(1,))
              encrypted1 = ArithmeticSharedTensor(tensor1)
              encrypted2 = ArithmeticSharedTensor(tensor2)
              reference = getattr(tensor1, func)(tensor2)
              encrypted_out = getattr(encrypted1, func)(encrypted2)
              self._check(
                  encrypted_out, reference, "private %s failed" % func
              )
  
          tensor = get_random_test_tensor(is_float=True)
          reference = tensor * tensor
          encrypted = ArithmeticSharedTensor(tensor)
          encrypted_out = encrypted.square()
          self._check(encrypted_out, reference, "square failed")
  
      # Test radd, rsub, and rmul
      reference = 2 + tensor1
      encrypted = ArithmeticSharedTensor(tensor1)
      encrypted_out = 2 + encrypted
      self._check(encrypted_out, reference, "right add failed")
  
      reference = 2 - tensor1
      encrypted_out = 2 - encrypted
      self._check(encrypted_out, reference, "right sub failed")
  
      reference = 2 * tensor1
      encrypted_out = 2 * encrypted
      self._check(encrypted_out, reference, "right mul failed")
  
imports: import itertools
  import logging
  import unittest
  import crypten
  import crypten.communicator
  import torch
  import torch.nn.functional
  from crypten.common.rng import generate_random_ring_element
  from crypten.common.tensor_types import is_float_tensor
  from crypten.common.util import count_wraps
  from crypten.mpc.primitives import ArithmeticSharedTensor
  from test.multiprocess_test_case import get_random_test_tensor, MultiProcessTestCase
label_desc: Use image node_id 5 for calling the TestArithmetic obj's underlying member method code with example usage: obj.test_arithmetic() without return types