class_docstr: This class tests all functions of the ArithmeticSharedTensor.
docstr: Tests where() conditional element selection
docstr_tok: ["Tests","where","(",")","conditional","element","selection"]
code: def test_where(self) -> None:
      """Tests where() conditional element selection"""
      sizes = [(10,), (5, 10), (1, 5, 10)]
      y_types = [lambda x: x, ArithmeticSharedTensor]
  
      for size, y_type in itertools.product(sizes, y_types):
          tensor1 = get_random_test_tensor(size=size, is_float=True)
          encrypted_tensor1 = ArithmeticSharedTensor(tensor1)
          tensor2 = get_random_test_tensor(size=size, is_float=True)
          encrypted_tensor2 = y_type(tensor2)
  
          condition_tensor = (
              get_random_test_tensor(
                  max_value=1, size=size, is_float=False
              )
              + 1
          )
          condition_encrypted = ArithmeticSharedTensor(condition_tensor)
          condition_bool = condition_tensor.bool()
  
          reference_out = tensor1.where(condition_bool, tensor2)
  
          encrypted_out = encrypted_tensor1.where(
              condition_bool, encrypted_tensor2
          )
          y_is_private = y_type == ArithmeticSharedTensor
          self._check(
              encrypted_out,
              reference_out,
              f"{'private' if y_is_private else 'public'} y "
              "where failed with public condition",
          )
  
          encrypted_out = encrypted_tensor1.where(
              condition_encrypted, encrypted_tensor2
          )
          self._check(
              encrypted_out,
              reference_out,
              f"{'private' if y_is_private else 'public'} y "
              "where failed with private condition",
          )
  
          # test scalar y
          scalar = get_random_test_tensor(
              max_value=0, size=[1], is_float=True
          )
          self._check(
              encrypted_tensor1.where(condition_bool, scalar),
              tensor1.where(condition_bool, scalar),
              "where failed against scalar y with public condition",
          )
  
          self._check(
              encrypted_tensor1.where(condition_encrypted, scalar),
              tensor1.where(condition_bool, scalar),
              "where failed against scalar y with private condition",
          )
  
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 36 for calling the TestArithmetic obj's underlying member method code with example usage: obj.test_where() without return types