class_docstr: This class tests all functions of the ArithmeticSharedTensor.
docstr: Test broadcast functionality.
docstr_tok: ["Test","broadcast","functionality","."]
code: def test_broadcast(self) -> None:
      """Test broadcast functionality."""
      arithmetic_functions = ["add", "sub", "mul", "div"]
      arithmetic_sizes = [
          (),
          (1,),
          (2,),
          (1, 1),
          (1, 2),
          (2, 1),
          (2, 2),
          (1, 1, 1),
          (1, 1, 2),
          (1, 2, 1),
          (2, 1, 1),
          (2, 2, 2),
          (1, 1, 1, 1),
          (1, 1, 1, 2),
          (1, 1, 2, 1),
          (1, 2, 1, 1),
          (2, 1, 1, 1),
          (2, 2, 2, 2),
      ]
      matmul_sizes = [(1, 1), (1, 5), (5, 1), (5, 5)]
      batch_dims = [(), (1,), (5,), (1, 1), (1, 5), (5, 5)]
  
      for tensor_type in [lambda x: x, ArithmeticSharedTensor]:
          for func in arithmetic_functions:
              for size1, size2 in itertools.combinations(
                  arithmetic_sizes, 2
              ):
                  tensor1 = get_random_test_tensor(
                      size=size1, is_float=True
                  )
                  tensor2 = get_random_test_tensor(
                      size=size2, is_float=True
                  )
  
                  # ArithmeticSharedTensors can't divide by negative
                  # private values - MPCTensor overrides this to allow negatives
                  # multiply denom by 10 to avoid division by small num
                  if (
                      func == "div"
                      and tensor_type == ArithmeticSharedTensor
                  ):
                      continue
  
                  encrypted1 = ArithmeticSharedTensor(tensor1)
                  encrypted2 = tensor_type(tensor2)
                  reference = getattr(tensor1, func)(tensor2)
                  encrypted_out = getattr(encrypted1, func)(encrypted2)
  
                  private = isinstance(
                      encrypted2, ArithmeticSharedTensor
                  )
                  self._check(
                      encrypted_out,
                      reference,
                      "%s %s broadcast failed"
                      % ("private" if private else "public", func),
                  )
  
          for size in matmul_sizes:
              for batch1, batch2 in itertools.combinations(
                  batch_dims, 2
              ):
                  size1 = (*batch1, *size)
                  size2 = (*batch2, *size)
  
                  tensor1 = get_random_test_tensor(
                      size=size1, is_float=True
                  )
                  tensor2 = get_random_test_tensor(
                      size=size2, is_float=True
                  )
                  tensor2 = tensor1.transpose(-2, -1)
  
                  encrypted1 = ArithmeticSharedTensor(tensor1)
                  encrypted2 = tensor_type(tensor2)
  
                  reference = tensor1.matmul(tensor2)
                  encrypted_out = encrypted1.matmul(encrypted2)
                  private = isinstance(
                      encrypted2, ArithmeticSharedTensor
                  )
                  self._check(
                      encrypted_out,
                      reference,
                      "%s matmul broadcast failed"
                      % ("private" if private else "public"),
                  )
  
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 31 for calling the TestArithmetic obj's underlying member method code with example usage: obj.test_broadcast() without return types