class_docstr: This class tests the crypten.nn package.
class_docstr_tok: ["This","class","tests","the","crypten.nn","package","."]
docstr: Tests crypten.nn.Graph module.
code: def test_graph(self):
      """
      Tests crypten.nn.Graph module.
      """
      for compute_gradients in [True, False]:
          for num_inputs in [1, 2]:
              # define test case:
              input_size = (3, 10)
              input = get_random_test_tensor(
                  size=input_size, is_float=True
              )
              input2 = get_random_test_tensor(
                  size=input_size, is_float=True
              )
              encr_input = crypten.cryptensor(input)
              encr_input2 = crypten.cryptensor(input2)
              encr_input.requires_grad = compute_gradients
              encr_input2.requires_grad = compute_gradients
  
              # for two inputs, sum the inputs first:
              if num_inputs == 1:
                  graph = crypten.nn.Graph("input", "output")
              elif num_inputs == 2:
                  graph = crypten.nn.Graph(
                      ["input1", "input2"], "output"
                  )
                  graph.add_module(
                      "input", crypten.nn.Add(), ["input1", "input2"]
                  )
              else:
                  raise ValueError(
                      f"Unsupported value of inputs: {num_inputs}"
                  )
  
              # test residual block with subsequent linear layer:
              linear1 = get_random_linear(input_size[1], input_size[1])
              linear2 = get_random_linear(input_size[1], input_size[1])
              graph.add_module(
                  "linear", linear_to_crypten(linear1), ["input"]
              )
              graph.add_module(
                  "residual", crypten.nn.Add(), ["input", "linear"]
              )
              graph.add_module(
                  "output", linear_to_crypten(linear2), ["residual"]
              )
              graph.encrypt()
  
              # check container:
              self.assertTrue(graph.encrypted, "nn.Graph not encrypted")
              for module in graph.modules():
                  self.assertTrue(
                      module.encrypted, "module not encrypted"
                  )
  
              # compare output to reference:
              if num_inputs == 1:
                  encr_output = graph(encr_input)
                  reference = linear2(linear1(input) + input)
              elif num_inputs == 2:
                  encr_output = graph(encr_input, encr_input2)
                  reference = linear2(
                      linear1(input + input2) + input + input2
                  )
              else:
                  raise ValueError(
                      f"Unsupported value of inputs: {num_inputs}"
                  )
              self._check(
                  encr_output, reference, "nn.Graph forward failed"
              )
  
imports: import copy
  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.config import cfg
  from crypten.encoder import FixedPointEncoder
  from test.multiprocess_test_case import get_random_linear, get_random_test_tensor, MultiProcessTestCase, onehot
label_desc: Use image node_id 15 for calling the TestNN obj's underlying member method code with example usage: obj.test_graph() without return types