code: r"""Importing this file must **not** initialize CUDA context. test_distributed
  relies on this assumption to properly run. This means that when this is imported
  no CUDA calls shall be made, including torch.cuda.device_count(), etc.
  
  torch.testing._internal.common_cuda.py can freely initialize CUDA context when imported.
  """
  
  import sys
  import os
  import platform
  import re
  import gc
  import types
  from functools import partial
  import inspect
  import io
  import argparse
  import unittest
  import warnings
  import random
  import contextlib
  import socket
  import subprocess
  import time
  from collections import OrderedDict
  from contextlib import contextmanager
  from functools import wraps
  from itertools import product
  from copy import deepcopy
  from numbers import Number
  import tempfile
  import json
  from urllib.request import urlopen
  import __main__
  import errno
  from typing import cast, Any, Iterable, Optional
  
  from torch.testing._internal import expecttest
  from torch.testing import _compare_tensors_internal, _compare_scalars_internal, _compare_return_type
  
  import torch
  import torch.cuda
  from torch._utils_internal import get_writable_path
  from torch._six import string_classes
  import torch.backends.cudnn
  import torch.backends.mkl
  from enum import Enum
  from torch.autograd import gradcheck
  from torch.autograd.gradcheck import gradgradcheck
  
  torch.backends.disable_global_flags()
  
  IS_SANDCASTLE = os.getenv('SANDCASTLE') == '1' or os.getenv('TW_JOB_USER') == 'sandcastle'
  
  class ProfilingMode(Enum):
      LEGACY = 1
      SIMPLE = 2
      PROFILING = 3
  
  def cppProfilingFlagsToProfilingMode():
      old_prof_exec_state = torch._C._jit_set_profiling_executor(True)
      old_prof_mode_state = torch._C._jit_set_profiling_mode(True)
      torch._C._jit_set_profiling_executor(old_prof_exec_state)
      torch._C._jit_set_profiling_mode(old_prof_mode_state)
  
      if old_prof_exec_state:
          if old_prof_mode_state:
              return ProfilingMode.PROFILING
          else:
              return ProfilingMode.SIMPLE
      else:
          return ProfilingMode.LEGACY
  
  @contextmanager
  def enable_profiling_mode_for_profiling_tests():
      if GRAPH_EXECUTOR == ProfilingMode.PROFILING:
          old_prof_exec_state = torch._C._jit_set_profiling_executor(True)
          old_prof_mode_state = torch._C._jit_set_profiling_mode(True)
      try:
          yield
      finally:
          if GRAPH_EXECUTOR == ProfilingMode.PROFILING:
              torch._C._jit_set_profiling_executor(old_prof_exec_state)
              torch._C._jit_set_profiling_mode(old_prof_mode_state)
  
  @contextmanager
  def enable_profiling_mode():
      old_prof_exec_state = torch._C._jit_set_profiling_executor(True)
      old_prof_mode_state = torch._C._jit_set_profiling_mode(True)
      try:
          yield
      finally:
          torch._C._jit_set_profiling_executor(old_prof_exec_state)
          torch._C._jit_set_profiling_mode(old_prof_mode_state)
  
  @contextmanager
  def num_profiled_runs(num_runs):
      old_num_runs = torch._C._jit_set_num_profiled_runs(num_runs)
      try:
          yield
      finally:
          torch._C._jit_set_num_profiled_runs(old_num_runs)
  
  func_call = torch._C.ScriptFunction.__call__
  meth_call = torch._C.ScriptMethod.__call__
  
  def prof_callable(callable, *args, **kwargs):
      if 'profile_and_replay' in kwargs:
          del kwargs['profile_and_replay']
          if GRAPH_EXECUTOR == ProfilingMode.PROFILING:
              with enable_profiling_mode_for_profiling_tests():
                  callable(*args, **kwargs)
                  return callable(*args, **kwargs)
  
      return callable(*args, **kwargs)
  
  def prof_func_call(*args, **kwargs):
      return prof_callable(func_call, *args, **kwargs)
  
  def prof_meth_call(*args, **kwargs):
      return prof_callable(meth_call, *args, **kwargs)
  
  torch._C.ScriptFunction.__call__ = prof_func_call
  torch._C.ScriptMethod.__call__ = prof_meth_call
  
  def _get_test_report_path():
      # allow users to override the test file location. We need this
      # because the distributed tests run the same test file multiple
      # times with different configurations.
      override = os.environ.get('TEST_REPORT_SOURCE_OVERRIDE')
      test_source = override if override is not None else 'python-unittest'
      return os.path.join('test-reports', test_source)
  
  
  parser = argparse.ArgumentParser(add_help=False)
  parser.add_argument('--subprocess', action='store_true',
                      help='whether to run each test in a subprocess')
  parser.add_argument('--seed', type=int, default=1234)
  parser.add_argument('--accept', action='store_true')
  parser.add_argument('--ge_config', type=str)
  parser.add_argument('--repeat', type=int, default=1)
  parser.add_argument('--test_bailouts', action='store_true')
  parser.add_argument('--save-xml', nargs='?', type=str,
                      const=_get_test_report_path(),
                      default=_get_test_report_path() if bool(os.environ.get('IN_CIRCLECI')) else None)
  parser.add_argument('--discover-tests', action='store_true')
  parser.add_argument('--log-suffix', type=str, default="")
  parser.add_argument('--run-parallel', type=int, default=1)
  
  args, remaining = parser.parse_known_args()
  if args.ge_config == 'legacy':
      GRAPH_EXECUTOR = ProfilingMode.LEGACY
  elif args.ge_config == 'profiling':
      GRAPH_EXECUTOR = ProfilingMode.PROFILING
  elif args.ge_config == 'simple':
      GRAPH_EXECUTOR = ProfilingMode.SIMPLE
  else:
      # infer flags based on the default settings
      GRAPH_EXECUTOR = cppProfilingFlagsToProfilingMode()
  
  
  LOG_SUFFIX = args.log_suffix
  RUN_PARALLEL = args.run_parallel
  TEST_BAILOUTS = args.test_bailouts
  TEST_DISCOVER = args.discover_tests
  TEST_IN_SUBPROCESS = args.subprocess
  TEST_SAVE_XML = args.save_xml
  REPEAT_COUNT = args.repeat
  SEED = args.seed
  if not expecttest.ACCEPT:
      expecttest.ACCEPT = args.accept
  UNITTEST_ARGS = [sys.argv[0]] + remaining
  torch.manual_seed(SEED)
  
  def wait_for_process(p):
      try:
          return p.wait()
      except KeyboardInterrupt:
          # Give `p` a chance to handle KeyboardInterrupt. Without this,
          # `pytest` can't print errors it collected so far upon KeyboardInterrupt.
          exit_status = p.wait(timeout=5)
          if exit_status is not None:
              return exit_status
          else:
              p.kill()
              raise
      except:  # noqa E722, copied from python core library
          p.kill()
          raise
      finally:
          # Always call p.wait() to ensure exit
          p.wait()
  
  def shell(command, cwd=None, env=None):
      sys.stdout.flush()
      sys.stderr.flush()
      # The following cool snippet is copied from Py3 core library subprocess.call
      # only the with
      #   1. `except KeyboardInterrupt` block added for SIGINT handling.
      #   2. In Py2, subprocess.Popen doesn't return a context manager, so we do
      #      `p.wait()` in a `final` block for the code to be portable.
      #
      # https://github.com/python/cpython/blob/71b6c1af727fbe13525fb734568057d78cea33f3/Lib/subprocess.py#L309-L323
      assert not isinstance(command, torch._six.string_classes), "Command to shell should be a list or tuple of tokens"
      p = subprocess.Popen(command, universal_newlines=True, cwd=cwd, env=env)
      return wait_for_process(p)
  
  
  # Used to run the same test with different tensor types
  def repeat_test_for_types(dtypes):
      def repeat_helper(f):
          @wraps(f)
          def call_helper(self, *args):
              for dtype in dtypes:
                  with TestCase.subTest(self, dtype=dtype):
                      f(self, *args, dtype=dtype)
  
          return call_helper
      return repeat_helper
  
  # Environment variable `IS_PYTORCH_CI` is set in `.jenkins/common.sh`.
  IS_PYTORCH_CI = bool(os.environ.get('IS_PYTORCH_CI'))
  
  
  def discover_test_cases_recursively(suite_or_case):
      if isinstance(suite_or_case, unittest.TestCase):
          return [suite_or_case]
      rc = []
      for element in suite_or_case:
          rc.extend(discover_test_cases_recursively(element))
      return rc
  
  def get_test_names(test_cases):
      return ['.'.join(case.id().split('.')[-2:]) for case in test_cases]
  
  def chunk_list(lst, nchunks):
      return [lst[i::nchunks] for i in range(nchunks)]
  
  
  def run_tests(argv=UNITTEST_ARGS):
      if TEST_DISCOVER:
          suite = unittest.TestLoader().loadTestsFromModule(__main__)
          test_cases = discover_test_cases_recursively(suite)
          for name in get_test_names(test_cases):
              print(name)
      elif TEST_IN_SUBPROCESS:
          suite = unittest.TestLoader().loadTestsFromModule(__main__)
          test_cases = discover_test_cases_recursively(suite)
          failed_tests = []
          for case in test_cases:
              test_case_full_name = case.id().split('.', 1)[1]
              exitcode = shell([sys.executable] + argv + [test_case_full_name])
              if exitcode != 0:
                  failed_tests.append(test_case_full_name)
  
          assert len(failed_tests) == 0, "{} unit test(s) failed:\n\t{}".format(
              len(failed_tests), '\n\t'.join(failed_tests))
      elif RUN_PARALLEL > 1:
          suite = unittest.TestLoader().loadTestsFromModule(__main__)
          test_cases = discover_test_cases_recursively(suite)
          test_batches = chunk_list(get_test_names(test_cases), RUN_PARALLEL)
          processes = []
          for i in range(RUN_PARALLEL):
              command = [sys.executable] + argv + ['--log-suffix=-shard-{}'.format(i + 1)] + test_batches[i]
              processes.append(subprocess.Popen(command, universal_newlines=True))
          failed = False
          for p in processes:
              failed |= wait_for_process(p) != 0
          assert not failed, "Some test shards have failed"
      elif TEST_SAVE_XML is not None:
          # import here so that non-CI doesn't need xmlrunner installed
          import xmlrunner
          test_report_path = TEST_SAVE_XML + LOG_SUFFIX
          os.makedirs(test_report_path, exist_ok=True)
          verbose = '--verbose' in argv or '-v' in argv
          if verbose:
              print('Test results will be stored in {}'.format(test_report_path))
          unittest.main(argv=argv, testRunner=xmlrunner.XMLTestRunner(output=test_report_path, verbosity=2 if verbose else 1))
      elif REPEAT_COUNT > 1:
          for _ in range(REPEAT_COUNT):
              if not unittest.main(exit=False, argv=argv).result.wasSuccessful():
                  sys.exit(-1)
      else:
          unittest.main(argv=argv)
  
  IS_WINDOWS = sys.platform == "win32"
  IS_MACOS = sys.platform == "darwin"
  IS_PPC = platform.machine() == "ppc64le"
  
  if IS_WINDOWS:
      @contextmanager
      def TemporaryFileName():
          # Ideally we would like to not have to manually delete the file, but NamedTemporaryFile
          # opens the file, and it cannot be opened multiple times in Windows. To support Windows,
          # close the file after creation and try to remove it manually
          f = tempfile.NamedTemporaryFile(delete=False)
          try:
              f.close()
              yield f.name
          finally:
              os.unlink(f.name)
  else:
      @contextmanager  # noqa: T484
      def TemporaryFileName():
          with tempfile.NamedTemporaryFile() as f:
              yield f.name
  
  
  def _check_module_exists(name):
      r"""Returns if a top-level module with :attr:`name` exists *without**
      importing it. This is generally safer than try-catch block around a
      `import X`. It avoids third party libraries breaking assumptions of some of
      our tests, e.g., setting multiprocessing start method when imported
      (see librosa/#747, torchvision/#544).
      """
      import importlib
      import importlib.util
      spec = importlib.util.find_spec(name)
      return spec is not None
  
  TEST_NUMPY = _check_module_exists('numpy')
  TEST_SCIPY = _check_module_exists('scipy')
  TEST_MKL = torch.backends.mkl.is_available()
  TEST_NUMBA = _check_module_exists('numba')
  
  TEST_DILL = _check_module_exists('dill')
  
  TEST_LIBROSA = _check_module_exists('librosa')
  
  # Python 2.7 doesn't have spawn
  NO_MULTIPROCESSING_SPAWN = os.environ.get('NO_MULTIPROCESSING_SPAWN', '0') == '1'
  TEST_WITH_ASAN = os.getenv('PYTORCH_TEST_WITH_ASAN', '0') == '1'
  TEST_WITH_TSAN = os.getenv('PYTORCH_TEST_WITH_TSAN', '0') == '1'
  TEST_WITH_UBSAN = os.getenv('PYTORCH_TEST_WITH_UBSAN', '0') == '1'
  TEST_WITH_ROCM = os.getenv('PYTORCH_TEST_WITH_ROCM', '0') == '1'
  # Enables tests that are slow to run (disabled by default)
  TEST_WITH_SLOW = os.getenv('PYTORCH_TEST_WITH_SLOW', '0') == '1'
  
  # Disables non-slow tests (these tests enabled by default)
  # This is usually used in conjunction with TEST_WITH_SLOW to
  # run *only* slow tests.  (I could have done an enum, but
  # it felt a little awkward.
  TEST_SKIP_FAST = os.getenv('PYTORCH_TEST_SKIP_FAST', '0') == '1'
  
  if TEST_NUMPY:
      import numpy as np
  
      # Dict of NumPy dtype -> torch dtype (when the correspondence exists)
      numpy_to_torch_dtype_dict = {
          np.bool       : torch.bool,
          np.uint8      : torch.uint8,
          np.int8       : torch.int8,
          np.int16      : torch.int16,
          np.int32      : torch.int32,
          np.int64      : torch.int64,
          np.float16    : torch.float16,
          np.float32    : torch.float32,
          np.float64    : torch.float64,
          np.complex64  : torch.complex64,
          np.complex128 : torch.complex128
      }
  
      # Dict of torch dtype -> NumPy dtype
      torch_to_numpy_dtype_dict = {value : key for (key, value) in numpy_to_torch_dtype_dict.items()}
  
  ALL_TENSORTYPES = [torch.float,
                     torch.double,
                     torch.half]
  
  # bfloat16 bringup is currently only available on ROCm
  # ALL_TENSORTYPES2 will eventually be unified with ALL_TENSORTYPES
  # when bfloat16 bringup is complete on all platforms
  if TEST_WITH_ROCM:
      ALL_TENSORTYPES2 = [torch.float,
                          torch.double,
                          torch.half,
                          torch.bfloat16]
  else:
      ALL_TENSORTYPES2 = ALL_TENSORTYPES
  
  def skipIfRocm(fn):
      @wraps(fn)
      def wrapper(*args, **kwargs):
          if TEST_WITH_ROCM:
              raise unittest.SkipTest("test doesn't currently work on the ROCm stack")
          else:
              fn(*args, **kwargs)
      return wrapper
  
  
  def skipIfCompiledWithoutNumpy(fn):
      # Even if the numpy module is present, if `USE_NUMPY=0` is used during the
      # build, numpy tests will fail
      numpy_support = TEST_NUMPY
      if numpy_support:
          try:
              # The numpy module is present, verify that PyTorch is compiled with
              # numpy support
              torch.from_numpy(np.array([2, 2]))
          except RuntimeError:
              numpy_support = False
  
      @wraps(fn)
      def wrapper(*args, **kwargs):
          if not numpy_support:
              raise unittest.SkipTest("PyTorch was compiled without numpy support")
          else:
              fn(*args, **kwargs)
      return wrapper
  
  def _test_function(fn, device):
      def run_test_function(self):
          return fn(self, device)
      return run_test_function
  
  
  def skipIfNoLapack(fn):
      @wraps(fn)
      def wrapper(*args, **kwargs):
          if not torch._C.has_lapack:
              raise unittest.SkipTest('PyTorch compiled without Lapack')
          else:
              fn(*args, **kwargs)
      return wrapper
  
  
  def skipIfNotRegistered(op_name, message):
      """Wraps the decorator to hide the import of the `core`.
  
      Args:
          op_name: Check if this op is registered in `core._REGISTERED_OPERATORS`.
          message: message to fail with.
  
      Usage:
          @skipIfNotRegistered('MyOp', 'MyOp is not linked!')
              This will check if 'MyOp' is in the caffe2.python.core
      """
      try:
          from caffe2.python import core
          skipper = unittest.skipIf(op_name not in core._REGISTERED_OPERATORS,
                                    message)
      except ImportError:
          skipper = unittest.skip("Cannot import `caffe2.python.core`")
      return skipper
  
  
  def skipIfNoSciPy(fn):
      @wraps(fn)
      def wrapper(*args, **kwargs):
          if not TEST_SCIPY:
              raise unittest.SkipTest("test require SciPy, but SciPy not found")
          else:
              fn(*args, **kwargs)
      return wrapper
  
  
  def slowTest(fn):
      @wraps(fn)
      def wrapper(*args, **kwargs):
          if not TEST_WITH_SLOW:
              raise unittest.SkipTest("test is slow; run with PYTORCH_TEST_WITH_SLOW to enable test")
          else:
              fn(*args, **kwargs)
      wrapper.__dict__['slow_test'] = True
      return wrapper
  
  
  def skipCUDAMemoryLeakCheckIf(condition):
      def dec(fn):
          if getattr(fn, '_do_cuda_memory_leak_check', True):  # if current True
              fn._do_cuda_memory_leak_check = not condition
          return fn
      return dec
  
  def skipCUDANonDefaultStreamIf(condition):
      def dec(fn):
          if getattr(fn, '_do_cuda_non_default_stream', True):  # if current True
              fn._do_cuda_non_default_stream = not condition
          return fn
      return dec
  
  def suppress_warnings(fn):
      @wraps(fn)
      def wrapper(*args, **kwargs):
          with warnings.catch_warnings():
              warnings.simplefilter("ignore")
              fn(*args, **kwargs)
      return wrapper
  
  
  def get_cpu_type(type_name):
      module, name = type_name.rsplit('.', 1)
      assert module == 'torch.cuda'
      return getattr(torch, name)
  
  
  def get_gpu_type(type_name):
      if isinstance(type_name, type):
          type_name = '{}.{}'.format(type_name.__module__, type_name.__name__)
      module, name = type_name.rsplit('.', 1)
      assert module == 'torch'
      return getattr(torch.cuda, name)
  
  
  def to_gpu(obj, type_map=None):
      if type_map is None:
          type_map = {}
      if isinstance(obj, torch.Tensor):
          assert obj.is_leaf
          t = type_map.get(obj.type(), get_gpu_type(obj.type()))
          with torch.no_grad():
              res = obj.clone().type(t)
              res.requires_grad = obj.requires_grad
          return res
      elif torch.is_storage(obj):
          return obj.new().resize_(obj.size()).copy_(obj)
      elif isinstance(obj, list):
          return [to_gpu(o, type_map) for o in obj]
      elif isinstance(obj, tuple):
          return tuple(to_gpu(o, type_map) for o in obj)
      else:
          return deepcopy(obj)
  
  
  def get_function_arglist(func):
      return inspect.getfullargspec(func).args
  
  
  def set_rng_seed(seed):
      torch.manual_seed(seed)
      random.seed(seed)
      if TEST_NUMPY:
          np.random.seed(seed)
  
  
  @contextlib.contextmanager
  def freeze_rng_state():
      rng_state = torch.get_rng_state()
      if torch.cuda.is_available():
          cuda_rng_state = torch.cuda.get_rng_state()
      yield
      if torch.cuda.is_available():
          torch.cuda.set_rng_state(cuda_rng_state)
      torch.set_rng_state(rng_state)
  
  @contextlib.contextmanager
  def set_default_dtype(dtype):
      saved_dtype = torch.get_default_dtype()
      torch.set_default_dtype(dtype)
      yield
      torch.set_default_dtype(saved_dtype)
  
  def iter_indices(tensor):
      if tensor.dim() == 0:
          return range(0)
      if tensor.dim() == 1:
          return range(tensor.size(0))
      return product(*(range(s) for s in tensor.size()))
  
  
  def is_iterable(obj):
      try:
          iter(obj)
          return True
      except TypeError:
          return False
  
  class CudaNonDefaultStream():
      def __enter__(self):
          # Before starting CUDA test save currently active streams on all
          # CUDA devices and set new non default streams to all CUDA devices
          # to ensure CUDA tests do not use default stream by mistake.
          beforeDevice = torch.cuda.current_device()
          self.beforeStreams = []
          for d in range(torch.cuda.device_count()):
              self.beforeStreams.append(torch.cuda.current_stream(d))
              deviceStream = torch.cuda.Stream(device=d)
              torch._C._cuda_setStream(deviceStream._cdata)
          torch._C._cuda_setDevice(beforeDevice)
  
      def __exit__(self, exec_type, exec_value, traceback):
          # After completing CUDA test load previously active streams on all
          # CUDA devices.
          beforeDevice = torch.cuda.current_device()
          for d in range(torch.cuda.device_count()):
              torch._C._cuda_setStream(self.beforeStreams[d]._cdata)
          torch._C._cuda_setDevice(beforeDevice)
  
  class CudaMemoryLeakCheck():
      def __init__(self, testcase, name=None):
          self.name = testcase.id() if name is None else name
          self.testcase = testcase
  
          # initialize context & RNG to prevent false positive detections
          # when the test is the first to initialize those
          from torch.testing._internal.common_cuda import initialize_cuda_context_rng
          initialize_cuda_context_rng()
  
      @staticmethod
      def get_cuda_memory_usage():
          # we don't need CUDA synchronize because the statistics are not tracked at
          # actual freeing, but at when marking the block as free.
          num_devices = torch.cuda.device_count()
          gc.collect()
          return tuple(torch.cuda.memory_allocated(i) for i in range(num_devices))
  
      def __enter__(self):
          self.befores = self.get_cuda_memory_usage()
  
      def __exit__(self, exec_type, exec_value, traceback):
          # Don't check for leaks if an exception was thrown
          if exec_type is not None:
              return
  
          afters = self.get_cuda_memory_usage()
  
          for i, (before, after) in enumerate(zip(self.befores, afters)):
              self.testcase.assertEqual(
                  before, after, msg='{} leaked {} bytes CUDA memory on device {}'.format(
                      self.name, after - before, i))
  
  #  "min_satisfying_examples" setting has been deprecated in hypythesis
  #  3.56.0 and removed in hypothesis 4.x
  try:
      import hypothesis
  
      def settings(*args, **kwargs):
          if 'min_satisfying_examples' in kwargs and hypothesis.version.__version_info__ >= (3, 56, 0):
              kwargs.pop('min_satisfying_examples')
          return hypothesis.settings(*args, **kwargs)
  
  
      hypothesis.settings.register_profile(
          "pytorch_ci",
          settings(
              derandomize=True,
              suppress_health_check=[hypothesis.HealthCheck.too_slow],
              database=None,
              max_examples=100,
              verbosity=hypothesis.Verbosity.normal))
      hypothesis.settings.register_profile(
          "dev",
          settings(
              suppress_health_check=[hypothesis.HealthCheck.too_slow],
              database=None,
              max_examples=10,
              verbosity=hypothesis.Verbosity.normal))
      hypothesis.settings.register_profile(
          "debug",
          settings(
              suppress_health_check=[hypothesis.HealthCheck.too_slow],
              database=None,
              max_examples=1000,
              verbosity=hypothesis.Verbosity.verbose))
  
      hypothesis.settings.load_profile(
          "pytorch_ci" if IS_PYTORCH_CI else os.getenv('PYTORCH_HYPOTHESIS_PROFILE',
                                                       'dev')
      )
  except ImportError:
      print('Fail to import hypothesis in common_utils, tests are not derandomized')
  
  disabled_test_from_issues = None
  def check_disabled(test_name):
      global disabled_test_from_issues
      if disabled_test_from_issues is None:
          disabled_test_from_issues = {}
  
          def read_and_process():
              url = 'https://raw.githubusercontent.com/zdevito/pytorch_disabled_tests/master/result.json'
              contents = urlopen(url, timeout=1).read().decode('utf-8')
              the_response = json.loads(contents)
              for item in the_response['items']:
                  title = item['title']
                  key = 'DISABLED '
                  if title.startswith(key):
                      test_name = title[len(key):].strip()
                      disabled_test_from_issues[test_name] = item['html_url']
  
          if not IS_SANDCASTLE and os.getenv("PYTORCH_RUN_DISABLED_TESTS", "0") != "1":
              try:
                  read_and_process()
              except Exception:
                  print("Couldn't download test skip set, leaving all tests enabled...")
  
      if test_name in disabled_test_from_issues:
          raise unittest.SkipTest(
              "Test is disabled because an issue exists disabling it: {}".format(disabled_test_from_issues[test_name]) +
              " To enable set the environment variable PYTORCH_RUN_DISABLED_TESTS=1")
  
  # Acquires the comparison dtype, required since isclose
  # requires both inputs have the same dtype, and isclose is not supported
  # for some device x dtype combinations.
  # NOTE: Remaps bfloat16 to float32 since neither the CPU or CUDA device types
  #  support needed bfloat16 comparison methods.
  # NOTE: Remaps float16 to float32 on CPU since the CPU device type doesn't
  #   support needed float16 comparison methods.
  # TODO: Update this once bfloat16 and float16 are better supported.
  def get_comparison_dtype(a, b):
      # TODO: update this when promote_types supports bfloat16 and/or
      # isclose supports bfloat16.
      a_dtype = torch.float32 if a.dtype is torch.bfloat16 else a.dtype
      b_dtype = torch.float32 if b.dtype is torch.bfloat16 else b.dtype
  
      compare_dtype = torch.promote_types(a_dtype, b_dtype)
  
      # non-CUDA (CPU, for example) float16 -> float32
      # TODO: update this when isclose is implemented for CPU float16
      if (compare_dtype is torch.float16 and
          (a.device != b.device or a.device.type != 'cuda' or
              b.device.type != 'cuda')):
          compare_dtype = torch.float32
  
      return compare_dtype
  
  class TestCase(expecttest.TestCase):
      # NOTE: "precision" lets classes and generated tests set minimum
      # atol values when comparing tensors. Used by @precisionOverride, for
      # example.
      # TODO: provide a better mechanism for generated tests to set rtol/atol.
      _precision: float = 0
  
      @property
      def precision(self) -> float:
          return self._precision
  
      @precision.setter
      def precision(self, prec: float) -> None:
          self._precision = prec
  
      _do_cuda_memory_leak_check = False
      _do_cuda_non_default_stream = False
  
      def __init__(self, method_name='runTest'):
          super().__init__(method_name)
  
          test_method = getattr(self, method_name, None)
          if test_method is not None:
              # Wraps the tested method if we should do CUDA memory check.
              self._do_cuda_memory_leak_check &= getattr(test_method, '_do_cuda_memory_leak_check', True)
              # FIXME: figure out the flaky -1024 anti-leaks on windows. See #8044
              if self._do_cuda_memory_leak_check and not IS_WINDOWS:
                  self.wrap_with_cuda_policy(method_name, self.assertLeaksNoCudaTensors)
  
              # Wraps the tested method if we should enforce non default CUDA stream.
              self._do_cuda_non_default_stream &= getattr(test_method, '_do_cuda_non_default_stream', True)
              if self._do_cuda_non_default_stream and not IS_WINDOWS and not TEST_WITH_ROCM:
                  self.wrap_with_cuda_policy(method_name, self.enforceNonDefaultStream)
  
      def assertLeaksNoCudaTensors(self, name=None):
          name = self.id() if name is None else name
          return CudaMemoryLeakCheck(self, name)
  
      def enforceNonDefaultStream(self):
          return CudaNonDefaultStream()
  
      def wrap_with_cuda_policy(self, method_name, policy):
          test_method = getattr(self, method_name)
          # the import below may initialize CUDA context, so we do it only if
          # self._do_cuda_memory_leak_check or self._do_cuda_non_default_stream
          # is True.
          from torch.testing._internal.common_cuda import TEST_CUDA
          fullname = self.id().lower()  # class_name.method_name
          if TEST_CUDA and ('gpu' in fullname or 'cuda' in fullname):
              setattr(self, method_name, self.wrap_method_with_cuda_policy(test_method, policy))
  
      def wrap_method_with_cuda_policy(self, method, policy):
          # Assumes that `method` is the tested function in `self`.
          # NOTE: Python Exceptions (e.g., unittest.Skip) keeps objects in scope
          #       alive, so this cannot be done in setUp and tearDown because
          #       tearDown is run unconditionally no matter whether the test
          #       passes or not. For the same reason, we can't wrap the `method`
          #       call in try-finally and always do the check.
          @wraps(method)
          def wrapper(self, *args, **kwargs):
              with policy():
                  method(*args, **kwargs)
          return types.MethodType(wrapper, self)
  
      def wrap_with_cuda_memory_check(self, method):
          return self.wrap_method_with_cuda_policy(method, self.assertLeaksNoCudaTensors)
  
  
      def setUp(self):
  
  
          if TEST_SKIP_FAST:
              if not getattr(self, self._testMethodName).__dict__.get('slow_test', False):
                  raise unittest.SkipTest("test is fast; we disabled it with PYTORCH_TEST_SKIP_FAST")
          check_disabled(str(self))
  
          set_rng_seed(SEED)
  
      def genSparseTensor(self, size, sparse_dim, nnz, is_uncoalesced, device='cpu'):
          # Assert not given impossible combination, where the sparse dims have
          # empty numel, but nnz > 0 makes the indices containing values.
          assert all(size[d] > 0 for d in range(sparse_dim)) or nnz == 0, 'invalid arguments'
  
          v_size = [nnz] + list(size[sparse_dim:])
          v = torch.randn(*v_size, device=device)
          i = torch.rand(sparse_dim, nnz, device=device)
          i.mul_(torch.tensor(size[:sparse_dim]).unsqueeze(1).to(i))
          i = i.to(torch.long)
          if is_uncoalesced:
              v = torch.cat([v, torch.randn_like(v)], 0)
              i = torch.cat([i, i], 1)
  
          x = torch.sparse_coo_tensor(i, v, torch.Size(size))
  
          if not is_uncoalesced:
              x = x.coalesce()
          else:
              # FIXME: `x` is a sparse view of `v`. Currently rebase_history for
              #        sparse views is not implemented, so this workaround is
              #        needed for inplace operations done on `x`, e.g., copy_().
              #        Remove after implementing something equivalent to CopySlice
              #        for sparse views.
              # NOTE: We do clone() after detach() here because we need to be able to change size/storage of x afterwards
              x = x.detach().clone()
          return x, x._indices().clone(), x._values().clone()
  
      def safeToDense(self, t):
          r = self.safeCoalesce(t)
          return r.to_dense()
  
      def safeCoalesce(self, t):
          tc = t.coalesce()
          self.assertEqual(tc.to_dense(), t.to_dense())
          self.assertTrue(tc.is_coalesced())
  
          # Our code below doesn't work when nnz is 0, because
          # then it's a 0D tensor, not a 2D tensor.
          if t._nnz() == 0:
              self.assertEqual(t._indices(), tc._indices())
              self.assertEqual(t._values(), tc._values())
              return tc
  
          value_map = {}
          for idx, val in zip(t._indices().t(), t._values()):
              idx_tup = tuple(idx.tolist())
              if idx_tup in value_map:
                  value_map[idx_tup] += val
              else:
                  value_map[idx_tup] = val.clone() if isinstance(val, torch.Tensor) else val
  
          new_indices = sorted(list(value_map.keys()))
          new_values = [value_map[idx] for idx in new_indices]
          if t._values().ndimension() < 2:
              new_values = t._values().new(new_values)
          else:
              new_values = torch.stack(new_values)
  
          new_indices = t._indices().new(new_indices).t()
          tg = t.new(new_indices, new_values, t.size())
  
          self.assertEqual(tc._indices(), tg._indices())
          self.assertEqual(tc._values(), tg._values())
  
          if t.is_coalesced():
              self.assertEqual(tc._indices(), t._indices())
              self.assertEqual(tc._values(), t._values())
  
          return tg
  
      # Compares the given Torch and NumPy functions on the given tensor-like object.
      # NOTE: both torch_fn and np_fn should be functions that take a single
      #   tensor (array). If the torch and/or NumPy function require additional
      #   arguments then wrap the function in a lambda or pass a partial function.
      # TODO: support bfloat16 comparisons
      # TODO: add args/kwargs for passing to assertEqual (e.g. rtol, atol)
      def compare_with_numpy(self, torch_fn, np_fn, tensor_like, device=None, dtype=None):
          assert TEST_NUMPY
          assert dtype is not torch.bfloat16
  
          if isinstance(tensor_like, torch.Tensor):
              assert device is None
              assert dtype is None
              a = tensor_like.detach().cpu().numpy()
              t = tensor_like
          else:
              a = np.array(tensor_like, dtype=torch_to_numpy_dtype_dict[dtype])
              t = torch.tensor(tensor_like, device=device, dtype=dtype)
  
          np_result = np_fn(a)
          torch_result = torch_fn(t).cpu()
  
          # Converts arrays to tensors
          if isinstance(np_result, np.ndarray):
              try:
                  np_result = torch.from_numpy(np_result)
              except Exception:
                  # NOTE: copying an array before conversion is necessary when,
                  #   for example, the array has negative strides.
                  np_result = torch.from_numpy(np_result.copy())
  
          self.assertEqual(np_result, torch_result)
  
      # Some analysis of tolerance by logging tests from test_torch.py can be found
      # in https://github.com/pytorch/pytorch/pull/32538.
      # dtype name : (rtol, atol)
      dtype_precisions = {
          torch.float16    : (0.001, 1e-5),
          torch.bfloat16   : (0.016, 1e-5),
          torch.float32    : (1.3e-6, 1e-5),
          torch.float64    : (1e-7, 1e-7),
          torch.complex32  : (0.001, 1e-5),
          torch.complex64  : (1.3e-6, 1e-5),
          torch.complex128 : (1e-7, 1e-7),
      }
  
      # Returns the "default" rtol and atol for comparing scalars or
      # tensors of the given dtypes.
      def _getDefaultRtolAndAtol(self, dtype0, dtype1):
          rtol = max(self.dtype_precisions.get(dtype0, (0, 0))[0],
                     self.dtype_precisions.get(dtype1, (0, 0))[0])
          atol = max(self.dtype_precisions.get(dtype0, (0, 0))[1],
                     self.dtype_precisions.get(dtype1, (0, 0))[1])
  
          return rtol, atol
  
      # Checks if two dense tensors are equal(-ish), returning (True, None)
      #   when they are and (False, debug_msg) when they are not.
      # If exact_dtype is true both tensors must have the same dtype.
      # If exact_device is true both tensors must be on the same device.
      # See the "Test Framework Tensor 'Equality'" note for more details.
      # NOTE: tensors on different devices are moved to the CPU to be compared when
      #   exact_device is False.
      # NOTE: this function checks the tensors' devices, sizes, and dtypes
      #  and acquires the appropriate device, dtype, rtol and atol to compare
      #  them with. It then calls _compare_tensors_internal.
      def _compareTensors(self, a, b, *, rtol: Optional[float] = None, atol=None, equal_nan=True,
                          exact_dtype=True, exact_device=False) -> _compare_return_type:
          assert (atol is None) == (rtol is None)
          if not isinstance(a, torch.Tensor):
              return (False, "argument a, {0}, to _compareTensors is not a tensor!".format(a))
          if not isinstance(b, torch.Tensor):
              return (False, "argument b, {0}, to _compareTensors is not a tensor!".format(b))
  
          # Validates tensors are on the same device
          if exact_device and a.device != b.device:
              return (False, ("Attempted to compare equality of tensors on "
                              "different devices! Got devices {0} and "
                              "{1}.".format(a.device, b.device)))
  
          # Compares tensors of different devices on the CPU
          if a.device != b.device:
              a = a.cpu()
              b = b.cpu()
  
          # Checks size matches
          if a.size() != b.size():
              return (False, ("Attempted to compare equality of tensors with "
                              "different sizes. Got sizes {0} and {1}.").format(a.size(), b.size()))
  
          # Checks dtype (if exact_dtype)
          if exact_dtype and a.dtype is not b.dtype:
              return (False, ("Attempted to compare equality of tensors with "
                              "different dtypes. Got dtypes {0} and {1}.").format(a.dtype, b.dtype))
  
          # Acquires rtol and atol
          if rtol is None:
              rtol, atol = self._getDefaultRtolAndAtol(a.dtype, b.dtype)
  
          atol = max(atol, self.precision)
  
          # Converts to comparison dtype
          dtype = get_comparison_dtype(a, b)
          a = a.to(dtype)
          b = b.to(dtype)
  
          return _compare_tensors_internal(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
  
      # Checks if two scalars are equal(-ish), returning (True, None)
      #   when they are and (False, debug_msg) when they are not.
      # NOTE: this function just acquires rtol and atol
      #   before calling _compare_scalars_internal.
      def _compareScalars(self, a, b, *,
                          rtol: Optional[float] = None, atol: Optional[float] = None, equal_nan=True) -> _compare_return_type:
          # Acquires rtol and atol
          assert (atol is None) == (rtol is None)
          if rtol is None:
              if isinstance(a, complex) or isinstance(b, complex):
                  rtol, atol = self._getDefaultRtolAndAtol(torch.complex64, torch.complex64)
              elif isinstance(a, float) or isinstance(b, float):
                  rtol, atol = self._getDefaultRtolAndAtol(torch.float32, torch.float32)
              else:
                  rtol, atol = 0, 0
          atol = max(atol, self.precision)
  
          return _compare_scalars_internal(a, b, rtol=cast(float, rtol), atol=cast(float, atol), equal_nan=equal_nan)
  
      def assertEqualIgnoreType(self, *args, **kwargs) -> None:
          # If you are seeing this function used, that means test is written wrongly
          # and deserves detailed investigation
          return self.assertEqual(*args, exact_dtype=False, **kwargs)
  
      # Compares x and y
      # TODO: default exact_device to True
      def assertEqual(self, x, y, msg: Optional[str] = None, *,
                      atol: Optional[float] = None, rtol: Optional[float] = None,
                      equal_nan=True, exact_dtype=True, exact_device=False) -> None:
          assert (atol is None) == (rtol is None), "If one of atol or rtol is specified the other must be, too"
  
          # Tensor x Number and Number x Tensor comparisons
          if isinstance(x, torch.Tensor) and isinstance(y, Number):
              self.assertEqual(x.item(), y, atol=atol, rtol=rtol, msg=msg,
                               exact_dtype=exact_dtype, exact_device=exact_device)
          elif isinstance(y, torch.Tensor) and isinstance(x, Number):
              self.assertEqual(x, y.item(), atol=atol, rtol=rtol, msg=msg,
                               exact_dtype=exact_dtype, exact_device=exact_device)
          # Tensor x np.bool
          elif isinstance(x, torch.Tensor) and isinstance(y, np.bool_):
              self.assertEqual(x.item(), y, atol=atol, rtol=rtol, msg=msg,
                               exact_dtype=exact_dtype, exact_device=exact_device)
          elif isinstance(y, torch.Tensor) and isinstance(x, np.bool_):
              self.assertEqual(x, y.item(), atol=atol, rtol=rtol, msg=msg,
                               exact_dtype=exact_dtype, exact_device=exact_device)
          # Tensor x Tensor
          elif isinstance(x, torch.Tensor) and isinstance(y, torch.Tensor):
              super().assertEqual(x.is_sparse, y.is_sparse, msg=msg)
              super().assertEqual(x.is_quantized, y.is_quantized, msg=msg)
              if x.is_sparse:
                  x = self.safeCoalesce(x)
                  y = self.safeCoalesce(y)
                  indices_result, debug_msg = self._compareTensors(x._indices(), y._indices(),
                                                                   rtol=rtol, atol=atol,
                                                                   equal_nan=equal_nan, exact_dtype=exact_dtype,
                                                                   exact_device=exact_device)
  
                  if not indices_result and msg is None:
                      assert debug_msg is not None
                      msg = "Sparse tensor indices failed to compare as equal! " + debug_msg
                  self.assertTrue(indices_result, msg=msg)
  
                  values_result, debug_msg = self._compareTensors(x._values(), y._values(),
                                                                  rtol=rtol, atol=atol,
                                                                  equal_nan=equal_nan, exact_dtype=exact_dtype,
                                                                  exact_device=exact_device)
  
                  if not values_result and msg is None:
                      assert debug_msg is not None
                      msg = "Sparse tensor values failed to compare as equal! " + debug_msg
                  self.assertTrue(values_result, msg=msg)
              elif x.is_quantized and y.is_quantized:
                  self.assertEqual(x.qscheme(), y.qscheme(), atol=atol, rtol=rtol,
                                   msg=msg, exact_dtype=exact_dtype,
                                   exact_device=exact_device)
  
                  if x.qscheme() == torch.per_tensor_affine:
                      self.assertEqual(x.q_scale(), y.q_scale(), atol=atol, rtol=rtol,
                                       msg=msg, exact_dtype=exact_dtype,
                                       exact_device=exact_device)
                      self.assertEqual(x.q_zero_point(), y.q_zero_point(),
                                       atol=atol, rtol=rtol, msg=msg,
                                       exact_dtype=exact_dtype, exact_device=exact_device)
                  elif x.qscheme() == torch.per_channel_affine:
                      self.assertEqual(x.q_per_channel_scales(), y.q_per_channel_scales(), atol=atol, rtol=rtol,
                                       msg=msg, exact_dtype=exact_dtype,
                                       exact_device=exact_device)
                      self.assertEqual(x.q_per_channel_zero_points(), y.q_per_channel_zero_points(),
                                       atol=atol, rtol=rtol, msg=msg,
                                       exact_dtype=exact_dtype, exact_device=exact_device)
                      self.assertEqual(x.q_per_channel_axis(), y.q_per_channel_axis(),
                                       atol=atol, rtol=rtol, msg=msg,
                                       exact_dtype=exact_dtype, exact_device=exact_device)
  
                  result, debug_msg = self._compareTensors(x.int_repr().to(torch.int32),
                                                           y.int_repr().to(torch.int32),
                                                           atol=atol, rtol=rtol,
                                                           exact_dtype=exact_dtype,
                                                           exact_device=exact_device)
  
                  if not result and msg is None:
                      assert debug_msg is not None
                      msg = "Quantized representations failed to compare as equal! " + debug_msg
                  self.assertTrue(result, msg=msg)
              else:
                  result, debug_msg = self._compareTensors(x, y, rtol=rtol, atol=atol,
                                                           equal_nan=equal_nan, exact_dtype=exact_dtype,
                                                           exact_device=exact_device)
  
                  if not result and msg is None:
                      assert debug_msg is not None
                      msg = "Tensors failed to compare as equal! " + debug_msg
                  self.assertTrue(result, msg=msg)
          elif isinstance(x, string_classes) and isinstance(y, string_classes):
              super().assertEqual(x, y, msg=msg)
          elif type(x) == set and type(y) == set:
              super().assertEqual(x, y, msg=msg)
          elif isinstance(x, dict) and isinstance(y, dict):
              if isinstance(x, OrderedDict) and isinstance(y, OrderedDict):
                  self.assertEqual(x.items(), y.items(), atol=atol, rtol=rtol,
                                   msg=msg, exact_dtype=exact_dtype,
                                   exact_device=exact_device)
              else:
                  self.assertEqual(set(x.keys()), set(y.keys()), atol=atol, rtol=rtol,
                                   msg=msg, exact_dtype=exact_dtype,
                                   exact_device=exact_device)
                  key_list = list(x.keys())
                  self.assertEqual([x[k] for k in key_list],
                                   [y[k] for k in key_list],
                                   atol=atol, rtol=rtol, msg=msg,
                                   exact_dtype=exact_dtype, exact_device=exact_device)
          elif isinstance(x, type) and isinstance(y, type):
              # See TestTorch.test_assert_equal_generic_meta
              super().assertEqual(x, y, msg=msg)
          elif is_iterable(x) and is_iterable(y):
              super().assertEqual(len(x), len(y), msg=msg)
              for x_, y_ in zip(x, y):
                  self.assertEqual(x_, y_, atol=atol, rtol=rtol, msg=msg,
                                   exact_dtype=exact_dtype, exact_device=exact_device)
          elif isinstance(x, bool) and isinstance(y, bool):
              self.assertTrue(x == y, msg=msg)
  
          # Scalar x Scalar
          elif isinstance(x, Number) and isinstance(y, Number):
              result, debug_msg = self._compareScalars(x, y, rtol=rtol, atol=atol,
                                                       equal_nan=equal_nan)
              if not result and msg is None:
                  assert debug_msg is not None
                  msg = "Scalars failed to compare as equal! " + debug_msg
              self.assertTrue(result, msg=msg)
          else:
              super().assertEqual(x, y, msg=msg)
  
      def assertAlmostEqual(self, x, y, *, places=None, msg=None, delta=None):
          prec = delta
          if places:
              prec = 10**(-places)
          rtol = None if prec is None else 0
          self.assertEqual(x, y, msg=msg, atol=prec, rtol=rtol)
  
      def assertNotEqual(self, x, y, msg: Optional[str] = None, *,
                         atol: Optional[float] = None, rtol: Optional[float] = None, **kwargs) -> None:
          with self.assertRaises(AssertionError, msg=msg):
              self.assertEqual(x, y, msg, atol=atol, rtol=rtol, **kwargs)
  
      def assertEqualTypeString(self, x, y) -> None:
          # This API is used simulate deprecated x.type() == y.type()
          self.assertEqual(x.device, y.device)
          self.assertEqual(x.dtype, y.dtype)
          self.assertEqual(x.is_sparse, y.is_sparse)
  
      def assertObjectIn(self, obj: Any, iterable: Iterable[Any]) -> None:
          for elem in iterable:
              if id(obj) == id(elem):
                  return
          raise AssertionError("object not found in iterable")
  
      # TODO: Support context manager interface
      # NB: The kwargs forwarding to callable robs the 'subname' parameter.
      # If you need it, manually apply your callable in a lambda instead.
      def assertExpectedRaises(self, exc_type, callable, *args, **kwargs):
          subname = None
          if 'subname' in kwargs:
              subname = kwargs['subname']
              del kwargs['subname']
          try:
              callable(*args, **kwargs)
          except exc_type as e:
              self.assertExpected(str(e), subname)
              return
          # Don't put this in the try block; the AssertionError will catch it
          self.fail(msg="Did not raise when expected to")
  
      def assertNotWarn(self, callable, msg=''):
          r"""
          Test if :attr:`callable` does not raise a warning.
          """
          with warnings.catch_warnings(record=True) as ws:
              warnings.simplefilter("always")  # allow any warning to be raised
              callable()
              self.assertTrue(len(ws) == 0, msg)
  
      @contextmanager
      def maybeWarnsRegex(self, category, regex=''):
          """Context manager for code that *may* warn, e.g. ``TORCH_WARN_ONCE``.
  
          This filters expected warnings from the test log and fails the test if
          any unexpected warnings are caught.
          """
          with warnings.catch_warnings(record=True) as ws:
              warnings.simplefilter("always")  # allow any warning to be raised
              # Ignore expected warnings
              warnings.filterwarnings("ignore", message=regex, category=category)
              try:
                  yield
              finally:
                  if len(ws) != 0:
                      msg = 'Caught unexpected warnings:\n'
                      for w in ws:
                          msg += warnings.formatwarning(
                              w.message, w.category, w.filename, w.lineno, w.line)
                          msg += '\n'
                      self.fail(msg)
  
      def assertExpected(self, s, subname=None):
          r"""
          Test that a string matches the recorded contents of a file
          derived from the name of this test and subname.  This file
          is placed in the 'expect' directory in the same directory
          as the test script. You can automatically update the recorded test
          output using --accept.
  
          If you call this multiple times in a single function, you must
          give a unique subname each time.
          """
          if not isinstance(s, str):
              raise TypeError("assertExpected is strings only")
  
          def remove_prefix(text, prefix):
              if text.startswith(prefix):
                  return text[len(prefix):]
              return text
          # NB: we take __file__ from the module that defined the test
          # class, so we place the expect directory where the test script
          # lives, NOT where test/common_utils.py lives.  This doesn't matter in
          # PyTorch where all test scripts are in the same directory as
          # test/common_utils.py, but it matters in onnx-pytorch
          module_id = self.__class__.__module__
          munged_id = remove_prefix(self.id(), module_id + ".")
          test_file = os.path.realpath(sys.modules[module_id].__file__)
          expected_file = os.path.join(os.path.dirname(test_file),
                                       "expect",
                                       munged_id)
  
          subname_output = ""
          if subname:
              expected_file += "-" + subname
              subname_output = " ({})".format(subname)
          expected_file += ".expect"
          expected = None
  
          def accept_output(update_type):
              print("Accepting {} for {}{}:\n\n{}".format(update_type, munged_id, subname_output, s))
              with open(expected_file, 'w') as f:
                  f.write(s)
  
          try:
              with open(expected_file) as f:
                  expected = f.read()
          except IOError as e:
              if e.errno != errno.ENOENT:
                  raise
              elif expecttest.ACCEPT:
                  return accept_output("output")
              else:
                  raise RuntimeError(
                      ("I got this output for {}{}:\n\n{}\n\n"
                       "No expect file exists; to accept the current output, run:\n"
                       "python {} {} --accept").format(munged_id, subname_output, s, __main__.__file__, munged_id))
  
          # a hack for JIT tests
          if IS_WINDOWS:
              expected = re.sub(r'CppOp\[(.+?)\]', 'CppOp[]', expected)
              s = re.sub(r'CppOp\[(.+?)\]', 'CppOp[]', s)
  
          # Adjust for producer_version
          expected = expected.replace(
              'producer_version: "XXX"',
              'producer_version: "{}"'.format(torch.onnx.producer_version)
          )
          if expecttest.ACCEPT:
              if expected != s:
                  return accept_output("updated output")
          else:
              if hasattr(self, "assertMultiLineEqual"):
                  # Python 2.7 only
                  # NB: Python considers lhs "old" and rhs "new".
                  self.assertMultiLineEqual(expected, s)
              else:
                  self.assertEqual(s, expected)
  
      def assertExpectedStripMangled(self, s, subname=None):
          s = re.sub(r'__torch__[^ ]+', '', s)
          self.assertExpected(s, subname)
  
      # returns captured stderr
      @staticmethod
      def runWithPytorchAPIUsageStderr(code):
          import subprocess
  
          env = os.environ.copy()
          env["PYTORCH_API_USAGE_STDERR"] = "1"
          pipes = subprocess.Popen(
              [sys.executable, '-c', code],
              stdout=subprocess.PIPE,
              stderr=subprocess.PIPE,
              env=env)
          return pipes.communicate()[1].decode('ascii')
  
      if sys.version_info < (3, 2):
          # assertRegexpMatches renamed to assertRegex in 3.2
          assertRegex = unittest.TestCase.assertRegexpMatches
          # assertRaisesRegexp renamed to assertRaisesRegex in 3.2
          assertRaisesRegex = unittest.TestCase.assertRaisesRegexp
  
      if sys.version_info < (3, 5):
          # assertNotRegexpMatches renamed to assertNotRegex in 3.5
          assertNotRegex = unittest.TestCase.assertNotRegexpMatches
  
  
  def download_file(url, binary=True):
      from urllib.parse import urlsplit
      from urllib import request, error
  
      filename = os.path.basename(urlsplit(url)[2])
      data_dir = get_writable_path(os.path.join(os.path.dirname(__file__), 'data'))
      path = os.path.join(data_dir, filename)
  
      if os.path.exists(path):
          return path
      try:
          data = request.urlopen(url, timeout=15).read()
          with open(path, 'wb' if binary else 'w') as f:
              f.write(data)
          return path
      except error.URLError:
          msg = "could not download test file '{}'".format(url)
          warnings.warn(msg, RuntimeWarning)
          raise unittest.SkipTest(msg)
  
  
  def find_free_port():
      sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
      sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
      sock.bind(('localhost', 0))
      sockname = sock.getsockname()
      sock.close()
      return sockname[1]
  
  # Errors that we can get in c10d initialization for which we should retry tests for.
  ADDRESS_IN_USE = "Address already in use"
  CONNECT_TIMEOUT = "connect() timed out."
  
  def retry_on_connect_failures(func=None, connect_errors=(ADDRESS_IN_USE)):
      """Reruns a test if the test returns a RuntimeError and the exception
      matches exactly with one of the strings in connect_errors."""
      # This if block is executed when using this function as a decorator with arguments.
      if func is None:
          return partial(retry_on_connect_failures, connect_errors=connect_errors)
  
      @wraps(func)
      def wrapper(*args, **kwargs):
          tries_remaining = 10
          while True:
              try:
                  return func(*args, **kwargs)
              except RuntimeError as error:
                  if str(error) in connect_errors:
                      tries_remaining -= 1
                      if tries_remaining == 0:
                          raise
                      time.sleep(random.random())
                      continue
                  raise
      return wrapper
  
  
  # Decorator to retry upon certain Exceptions.
  def retry(ExceptionToCheck, tries=3, delay=3, skip_after_retries=False):
      def deco_retry(f):
          @wraps(f)
          def f_retry(*args, **kwargs):
              mtries, mdelay = tries, delay
              while mtries > 1:
                  try:
                      return f(*args, **kwargs)
                  except ExceptionToCheck as e:
                      msg = "%s, Retrying in %d seconds..." % (str(e), mdelay)
                      print(msg)
                      time.sleep(mdelay)
                      mtries -= 1
              try:
                  return f(*args, **kwargs)
              except ExceptionToCheck as e:
                  raise unittest.SkipTest(f"Skipping after {tries} consecutive {str(e)}") from e if skip_after_retries else e
          return f_retry  # true decorator
      return deco_retry
  
  
  # Methods for matrix generation
  # Used in test_autograd.py and test_torch.py
  def prod_single_zero(dim_size):
      result = torch.randn(dim_size, dim_size)
      result[0, 1] = 0
      return result
  
  
  def random_square_matrix_of_rank(l, rank, dtype=torch.double, device='cpu'):
      assert rank <= l
      A = torch.randn(l, l, dtype=dtype, device=device)
      u, s, v = A.svd()
      for i in range(l):
          if i >= rank:
              s[i] = 0
          elif s[i] == 0:
              s[i] = 1
      return u.mm(torch.diag(s)).mm(v.transpose(0, 1))
  
  
  def random_symmetric_matrix(l, *batches, **kwargs):
      dtype = kwargs.get('dtype', torch.double)
      device = kwargs.get('device', 'cpu')
      A = torch.randn(*(batches + (l, l)), dtype=dtype, device=device)
      A = (A + A.transpose(-2, -1)).div_(2)
      return A
  
  
  def random_symmetric_psd_matrix(l, *batches, **kwargs):
      dtype = kwargs.get('dtype', torch.double)
      device = kwargs.get('device', 'cpu')
      A = torch.randn(*(batches + (l, l)), dtype=dtype, device=device)
      return torch.matmul(A, A.transpose(-2, -1))
  
  
  def random_symmetric_pd_matrix(matrix_size, *batch_dims, **kwargs):
      dtype = kwargs.get('dtype', torch.double)
      device = kwargs.get('device', 'cpu')
      A = torch.randn(*(batch_dims + (matrix_size, matrix_size)),
                      dtype=dtype, device=device)
      return torch.matmul(A, A.transpose(-2, -1)) \
          + torch.eye(matrix_size, dtype=dtype, device=device) * 1e-5
  
  
  def make_nonzero_det(A, sign=None, min_singular_value=0.1):
      u, s, v = A.svd()
      s.clamp_(min=min_singular_value)
      A = torch.matmul(u, torch.matmul(torch.diag_embed(s), v.transpose(-2, -1)))
      det = A.det()
      if sign is not None:
          if A.dim() == 2:
              det = det.item()
              if (det < 0) ^ (sign < 0):
                  A[0, :].neg_()
          else:
              cond = ((det < 0) ^ (sign < 0)).nonzero()
              if cond.size(0) > 0:
                  for i in range(cond.size(0)):
                      A[list(cond[i])][0, :].neg_()
      return A
  
  
  def random_fullrank_matrix_distinct_singular_value(matrix_size, *batch_dims,
                                                     **kwargs):
      dtype = kwargs.get('dtype', torch.double)
      device = kwargs.get('device', 'cpu')
      silent = kwargs.get("silent", False)
      if silent and not torch._C.has_lapack:
          return torch.ones(matrix_size, matrix_size, dtype=dtype, device=device)
  
      A = torch.randn(batch_dims + (matrix_size, matrix_size), dtype=dtype, device=device)
      u, _, v = A.svd()
      s = torch.arange(1., matrix_size + 1, dtype=dtype, device=device).mul_(1.0 / (matrix_size + 1)).diag()
      return u.matmul(s.expand(batch_dims + (matrix_size, matrix_size)).matmul(v.transpose(-2, -1)))
  
  
  def random_matrix(rows, columns, *batch_dims, **kwargs):
      """Return rectangular matrix or batches of rectangular matrices.
  
      Parameters:
        dtype - the data type
        device - the device kind
        singular - when True, the output will be singular
      """
      dtype = kwargs.get('dtype', torch.double)
      device = kwargs.get('device', 'cpu')
      silent = kwargs.get("silent", False)
      singular = kwargs.get("singular", False)
      if silent and not torch._C.has_lapack:
          return torch.ones(rows, columns, dtype=dtype, device=device)
  
      A = torch.randn(batch_dims + (rows, columns), dtype=dtype, device=device)
      u, _, v = A.svd(some=False)
      s = torch.zeros(rows, columns, dtype=dtype, device=device)
      k = min(rows, columns)
      for i in range(k):
          s[i, i] = float(i + 1) / (k + 1)
      if singular:
          # make matrix singular
          s[k - 1, k - 1] = 0
          if k > 2:
              # increase the order of singularity so that the pivoting
              # in LU factorization will be non-trivial
              s[0, 0] = 0
      return u.matmul(s.expand(batch_dims + (rows, columns)).matmul(v.transpose(-2, -1)))
  
  
  def random_lowrank_matrix(rank, rows, columns, *batch_dims, **kwargs):
      """Return rectangular matrix or batches of rectangular matrices with
      given rank.
      """
      B = random_matrix(rows, rank, *batch_dims, **kwargs)
      C = random_matrix(rank, columns, *batch_dims, **kwargs)
      return B.matmul(C)
  
  
  def random_sparse_matrix(rows, columns, density=0.01, **kwargs):
      """Return rectangular random sparse matrix within given density.
  
      The density of the result approaches to given density as the size
      of the matrix is increased and a relatively small value of density
      is specified but higher than min(rows, columns)/(rows * columns)
      for non-singular matrices.
      """
      dtype = kwargs.get('dtype', torch.double)
      device = kwargs.get('device', 'cpu')
      singular = kwargs.get("singular", False)
  
      k = min(rows, columns)
      nonzero_elements = max(min(rows, columns), int(rows * columns * density))
  
      row_indices = [i % rows for i in range(nonzero_elements)]
      column_indices = [i % columns for i in range(nonzero_elements)]
      random.shuffle(column_indices)
      indices = [row_indices, column_indices]
      values = torch.randn(nonzero_elements, dtype=dtype, device=device)
      # ensure that the diagonal dominates
      values *= torch.tensor([-float(i - j)**2 for i, j in zip(*indices)], dtype=dtype, device=device).exp()
      A = torch.sparse_coo_tensor(indices, values, (rows, columns), device=device)
      return A.coalesce()
  
  
  def random_sparse_pd_matrix(matrix_size, density=0.01, **kwargs):
      """Return random sparse positive-definite matrix with given density.
  
      The eigenvalues of the matrix are defined as::
        arange(1, matrix_size+1)/matrix_size
  
      Algorithm:
        A = diag(arange(1, matrix_size+1)/matrix_size)
        while <A density is smaller than required>:
            <choose random i, j in range(matrix_size), theta in [0, 2*pi]>
            R = <rotation matrix (i,j,theta)>
            A = R^T A R
      """
      import math
      torch = kwargs.get('torch', globals()['torch'])
      dtype = kwargs.get('dtype', torch.double)
      device = kwargs.get('device', 'cpu')
      data = dict([((i, i), float(i + 1) / matrix_size)
                   for i in range(matrix_size)])
  
  
      def multiply(data, N, i, j, cs, sn, left=True):
          for k in range(N):
              if left:
                  ik, jk = (k, i), (k, j)
              else:
                  ik, jk = (i, k), (j, k)
              aik, ajk = data.get(ik, 0), data.get(jk, 0)
              aik, ajk = cs * aik + sn * ajk, -sn * aik + cs * ajk
              if aik:
                  data[ik] = aik
              else:
                  data.pop(ik, None)
              if ajk:
                  data[jk] = ajk
              else:
                  data.pop(jk, None)
  
      target_nnz = density * matrix_size * matrix_size
      while len(data) < target_nnz:
          i = random.randint(0, matrix_size - 1)
          j = random.randint(0, matrix_size - 1)
          if i != j:
              theta = random.uniform(0, 2 * math.pi)
              cs = math.cos(theta)
              sn = math.sin(theta)
              multiply(data, matrix_size, i, j, cs, sn, left=True)
              multiply(data, matrix_size, i, j, cs, sn, left=False)
      icoords, jcoords, values = [], [], []
      for (i, j), v in sorted(data.items()):
          icoords.append(i)
          jcoords.append(j)
          values.append(v)
      indices = [icoords, jcoords]
      return torch.sparse_coo_tensor(indices, values, (matrix_size, matrix_size), dtype=dtype, device=device)
  
  
  def do_test_dtypes(self, dtypes, layout, device):
      for dtype in dtypes:
          if dtype != torch.float16:
              out = torch.zeros((2, 3), dtype=dtype, layout=layout, device=device)
              self.assertIs(dtype, out.dtype)
              self.assertIs(layout, out.layout)
              self.assertEqual(device, out.device)
  
  
  def do_test_empty_full(self, dtypes, layout, device):
      shape = torch.Size([2, 3])
  
      def check_value(tensor, dtype, layout, device, value, requires_grad):
          self.assertEqual(shape, tensor.shape)
          self.assertIs(dtype, tensor.dtype)
          self.assertIs(layout, tensor.layout)
          self.assertEqual(tensor.requires_grad, requires_grad)
          if tensor.is_cuda and device is not None:
              self.assertEqual(device, tensor.device)
          if value is not None:
              fill = tensor.new(shape).fill_(value)
              self.assertEqual(tensor, fill)
  
      def get_int64_dtype(dtype):
          module = '.'.join(str(dtype).split('.')[1:-1])
          if not module:
              return torch.int64
          return operator.attrgetter(module)(torch).int64
  
      default_dtype = torch.get_default_dtype()
      check_value(torch.empty(shape), default_dtype, torch.strided, -1, None, False)
      check_value(torch.full(shape, -5.), default_dtype, torch.strided, -1, None, False)
      for dtype in dtypes:
          for rg in {dtype.is_floating_point, False}:
              int64_dtype = get_int64_dtype(dtype)
              v = torch.empty(shape, dtype=dtype, device=device, layout=layout, requires_grad=rg)
              check_value(v, dtype, layout, device, None, rg)
              out = v.new()
              check_value(torch.empty(shape, out=out, device=device, layout=layout, requires_grad=rg),
                          dtype, layout, device, None, rg)
              check_value(v.new_empty(shape), dtype, layout, device, None, False)
              check_value(v.new_empty(shape, dtype=int64_dtype, device=device, requires_grad=False),
                          int64_dtype, layout, device, None, False)
              check_value(torch.empty_like(v), dtype, layout, device, None, False)
              check_value(torch.empty_like(v, dtype=int64_dtype, layout=layout, device=device, requires_grad=False),
                          int64_dtype, layout, device, None, False)
  
              if dtype is not torch.float16 and layout != torch.sparse_coo:
                  fv = 3
                  v = torch.full(shape, fv, dtype=dtype, layout=layout, device=device, requires_grad=rg)
                  check_value(v, dtype, layout, device, fv, rg)
                  check_value(v.new_full(shape, fv + 1), dtype, layout, device, fv + 1, False)
                  out = v.new()
                  check_value(torch.full(shape, fv + 2, out=out, device=device, layout=layout, requires_grad=rg),
                              dtype, layout, device, fv + 2, rg)
                  check_value(v.new_full(shape, fv + 3, dtype=int64_dtype, device=device, requires_grad=False),
                              int64_dtype, layout, device, fv + 3, False)
                  check_value(torch.full_like(v, fv + 4), dtype, layout, device, fv + 4, False)
                  check_value(torch.full_like(v, fv + 5,
                                              dtype=int64_dtype, layout=layout, device=device, requires_grad=False),
                              int64_dtype, layout, device, fv + 5, False)
  
  
  
  
  THESE_TAKE_WAY_TOO_LONG = {
      'test_Conv3d_groups',
      'test_conv_double_backward',
      'test_conv_double_backward_groups',
      'test_Conv3d_dilated',
      'test_Conv3d_stride_padding',
      'test_Conv3d_dilated_strided',
      'test_Conv3d',
      'test_Conv2d_dilated',
      'test_ConvTranspose3d_dilated',
      'test_ConvTranspose2d_dilated',
      'test_snli',
      'test_Conv2d',
      'test_Conv2d_padding',
      'test_ConvTranspose2d_no_bias',
      'test_ConvTranspose2d',
      'test_ConvTranspose3d',
      'test_Conv2d_no_bias',
      'test_matmul_4d_4d',
      'test_multinomial_invalid_probs',
  }
  
  
  running_script_path = None
  
  
  def set_running_script_path():
      global running_script_path
      try:
          running_file = os.path.abspath(os.path.realpath(sys.argv[0]))
          if running_file.endswith('.py'):  # skip if the running file is not a script
              running_script_path = running_file
      except Exception:
          pass
  
  
  def check_test_defined_in_running_script(test_case):
      if running_script_path is None:
          return
      test_case_class_file = os.path.abspath(os.path.realpath(inspect.getfile(test_case.__class__)))
      assert test_case_class_file == running_script_path, "Class of loaded TestCase \"{}\" " \
          "is not defined in the running script \"{}\", but in \"{}\". Did you " \
          "accidentally import a unittest.TestCase from another file?".format(
              test_case.id(), running_script_path, test_case_class_file)
  
  
  def load_tests(loader, tests, pattern):
      set_running_script_path()
      test_suite = unittest.TestSuite()
      for test_group in tests:
          for test in test_group:
              check_test_defined_in_running_script(test)
              test_suite.addTest(test)
      return test_suite
  
  
  class BytesIOContext(io.BytesIO):
      def __enter__(self):
          return self
  
      def __exit__(self, *args):
          pass
  
  def _assertGradAndGradgradChecks(test_case, apply_fn, inputs):
      # call assert function rather than returning a bool since it's nicer
      # if we get whether this failed on the gradcheck or the gradgradcheck.
      test_case.assertTrue(gradcheck(apply_fn, inputs))
      test_case.assertTrue(gradgradcheck(apply_fn, inputs))
  
  
  # Using @precisionOverride specific to your test is the recommended way
  # of doing this. These are just some values that worked for test_nn.
  dtype2prec_DONTUSE = {torch.float: 1e-5,
                        torch.double: 1e-5,
                        torch.half: 1e-2,
                        torch.bfloat16: 1e-1}
   #!/usr/bin/env python3
  
  from torch.testing._internal.distributed import ddp_under_dist_autograd_test
  from torch.testing._internal.common_utils import (
      run_tests,
  )
  
  class TestDdpUnderDistAutogradWrapper(ddp_under_dist_autograd_test.TestDdpUnderDistAutograd):
      pass
  
  class TestDdpComparison(ddp_under_dist_autograd_test.TestDdpComparison):
      pass
  
  if __name__ == "__main__":
      run_tests()
   from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from future.utils import viewkeys
  from multiprocessing import Process, Queue
  import numpy as np
  import os
  import shutil
  import tempfile
  import unittest
  import time
  from mock import Mock
  from hypothesis import assume, given
  import hypothesis.strategies as st
  
  from caffe2.proto import caffe2_pb2
  from caffe2.python import brew, core, cnn, data_parallel_model, dyndep, \
      model_helper, optimizer, rnn_cell, workspace
  from caffe2.python.test_util import TestCase
  
  
  dyndep.InitOpsLibrary("@/caffe2/caffe2/distributed:file_store_handler_ops")
  
  
  class TemporaryDirectory:
      def __enter__(self):
          self.tmpdir = tempfile.mkdtemp()
          return self.tmpdir
  
      def __exit__(self, type, value, traceback):
          shutil.rmtree(self.tmpdir)
  
  # Note(jiayq): we are yet to find out why Travis gives out an error in gloo
  # like:
  # RuntimeError: [enforce fail at /home/travis/build/caffe2/caffe2/third_party/gloo/gloo/transport/tcp/device.cc:113] ifa != nullptr. Unable to find interface for: [127.0.1.1]
  # See for example https://travis-ci.org/caffe2/caffe2/jobs/262433866
  # As a result, we will check if this is travis, and if yes, disable it.
  @unittest.skipIf(os.environ.get("TRAVIS"), "DPMTest has a known issue with Travis.")
  class DataParallelModelTest(TestCase):
  
      def run_model(self, devices, gpu):
          '''
          Helper function for test_equiv
          '''
          def input_builder_fun(model):
              return None
  
          def model_build_fun(model, loss_scale):
              fc = model.FC("data", "fc", 16, 1,
                            ("ConstantFill", {}), ("ConstantFill", {}))
              fc_fl = model.FlattenToVec(fc, "fc_fl")
              sigm = model.Sigmoid(fc_fl, "sigm")
              sq = model.SquaredL2Distance([sigm, "label"], "sq")
              loss = model.AveragedLoss(sq, "loss")
              loss = model.Scale(loss, scale=loss_scale)
  
              # For testing explicit sync
              model.param_init_net.UniformFill([], ["sync_num"], shape=[1])
              return [loss]
  
          def add_optimizer(model):
              return optimizer.build_sgd(
                  model,
                  0.1,
                  policy="fixed",
                  max_gradient_norm=5.0,
                  allow_lr_injection=True,
              )
  
          workspace.ResetWorkspace()
          model = cnn.CNNModelHelper(
              order="NHWC",
              name="test{}".format(devices),
          )
          data_parallel_model.Parallelize(
              model,
              input_builder_fun=input_builder_fun,
              forward_pass_builder_fun=model_build_fun,
              optimizer_builder_fun=add_optimizer,
              devices=devices,
              cpu_device=not gpu,
              shared_model=not gpu,
              combine_spatial_bn=not gpu,
          )
          data_parallel_model.AddBlobSync(model, ["sync_num"])
  
          # Light test for LR names
          lr_names = data_parallel_model.GetLearningRateBlobNames(model)
          self.assertGreater(len(lr_names), 0)
  
          np.random.seed(2603)
  
          # Each run has same input, independent of number of gpus
          batch_size = 64
          for i in range(0, 10):
              full_data = np.random.rand(batch_size, 16)
              full_labels = np.round(full_data[:, 0])
              batch_per_device = batch_size // len(devices)
  
              for (j, g) in enumerate(devices):
                  st = j * batch_per_device
                  en = st + batch_per_device
                  data = full_data[st:en, :].astype(np.float32)
                  labels = full_labels[st:en].astype(np.float32)
                  with core.DeviceScope(core.DeviceOption(model._device_type, g)):
                      workspace.FeedBlob(
                          "{}_{}/data".format(model._device_prefix, g), data
                      )
                      workspace.FeedBlob(
                          "{}_{}/label".format(model._device_prefix, g), labels
                      )
  
              if i == 0:
                  workspace.RunNetOnce(model.param_init_net)
                  workspace.CreateNet(model.net)
  
              workspace.FeedBlob(
                  model._device_prefix + "_0/sync_num",
                  np.array([i * 2]).astype(np.float32),
                  device_option=core.DeviceOption(model._device_type, 0))
              workspace.RunNet(model.net.Proto().name)
  
              # Test AddBlobSync
              for j in model._devices:
                  sync = workspace.FetchBlob(
                      model._device_prefix + "_{}/sync_num".format(j))[0]
                  self.assertTrue(abs(sync - i * 2) < 0.01)
  
          return workspace.FetchBlob("{}_0/fc_w".format(model._device_prefix))
  
      def run_test_locally(self, fn, device_option=None, **kwargs):
          # Queue for assertion errors on subprocesses
          queue = Queue()
  
          # Capture any exception thrown by the subprocess
          def run_fn(*args, **kwargs):
              try:
                  if device_option is None:
                      fn(*args, **kwargs)
                      workspace.ResetWorkspace()
                  else:
                      with core.DeviceScope(device_option):
                          fn(*args, **kwargs)
                          workspace.ResetWorkspace()
              except Exception as ex:
                  queue.put(ex)
  
          # Start N processes in the background
          procs = []
          for i in range(kwargs['comm_size']):
              kwargs['comm_rank'] = i
              proc = Process(
                  target=run_fn,
                  kwargs=kwargs)
              proc.start()
              procs.append(proc)
  
          # Test complete, join background processes
          while len(procs) > 0:
              proc = procs.pop(0)
              while proc.is_alive():
                  proc.join(1)
  
                  # Raise exception if we find any.
                  # Note that the following is executed ALSO after
                  # the last process was joined, so if ANY exception
                  # was raised, it will be re-raised here.
                  if not queue.empty():
                      raise queue.get()
  
      def test_equiv(self):
          '''
          Test that the model produces exactly same results given
          total batchsize, independent of number of GPUs.
          '''
          for gpu in [True, False]:
              if gpu and (not workspace.has_gpu_support or
                          workspace.NumCudaDevices() < 2):
                  continue
              result_2gpus = self.run_model([0, 1], gpu=gpu)
              result_1gpus = self.run_model([0], gpu=gpu)
  
              self.assertTrue(np.allclose(result_1gpus, result_2gpus))
  
              if not gpu or workspace.NumCudaDevices() >= 4:
                  result_4gpus = self.run_model(list(range(4)), gpu=gpu)
                  self.assertTrue(np.allclose(result_1gpus, result_4gpus))
  
              if not gpu or workspace.NumCudaDevices() >= 8:
                  result_8gpus = self.run_model(list(range(8)), gpu=gpu)
                  self.assertTrue(np.allclose(result_1gpus, result_8gpus))
  
              if not gpu or workspace.NumCudaDevices() >= 16:
                  result_16gpus = self.run_model(list(range(16)), gpu=gpu)
                  self.assertTrue(np.allclose(result_1gpus, result_16gpus))
  
      def test_checkpoint_params(self):
          def add_input_ops(model):
              pass
  
          def add_model_ops(model, loss_scale):
              model.NHWC2NCHW("data", "data_nchw")
              model.Conv("data_nchw", 'conv1', 3, 64,
                         weight_init=("MSRAFill", {}), kernel=7,
                         stride=2, pad=3, no_bias=0)
              model.SpatialBN('conv1', 'conv1_spatbn_relu', 64, epsilon=1e-3, is_test=False)
              model.Relu('conv1_spatbn_relu', 'conv1_spatbn_relu')
              model.MaxPool('conv1_spatbn_relu', 'pool1', kernel=3, stride=2)
              model.FC('pool1', 'fc', dim_in=(64 * 56 * 56), dim_out=100)
              model.Sigmoid('fc', 'fc_sigm')
              model.Softmax('fc_sigm', 'softmax')
              model.LabelCrossEntropy(['softmax', 'label'], 'xent')
              loss = model.AveragedLoss('xent', 'loss')
  
              # Add a duplicate param init to ensure it does not cause issues
              model.param_init_net.ConstantFill(
                  [], ["fc_w"], shape=((64 * 56 * 56), 1000)
              )
              return [loss]
  
          def add_optimizer(model):
              optimizer.build_sgd(model, 0.1, policy="fixed", momentum=0.9)
  
          model = cnn.CNNModelHelper(
              order="NHWC",
              name="test",
          )
          data_parallel_model.Parallelize_CPU(
              model,
              input_builder_fun=add_input_ops,
              forward_pass_builder_fun=add_model_ops,
              optimizer_builder_fun=add_optimizer,
              devices=[1, 2, 3],
          )
  
          # Only gpu_1 params should be returned (gpu_1 is the first gpu)
          checkpoint_params = data_parallel_model.GetCheckpointParams(model)
          for p in model.GetParams("cpu_1/"):
              self.assertTrue(p in checkpoint_params)
              self.assertTrue(p + "_momentum" in checkpoint_params)
          for p in model.GetParams("cpu_2/"):
              self.assertFalse(p in checkpoint_params)
          self.assertTrue(
              core.BlobReference("cpu_1/fc_w_momentum") in checkpoint_params)
          for c in model.GetComputedParams("cpu_1/"):
              self.assertTrue(c in checkpoint_params)
          for c in model.GetComputedParams("cpu_2/"):
              self.assertFalse(c in checkpoint_params)
          self.assertFalse(core.BlobReference("cpu_1/data") in checkpoint_params)
          self.assertTrue(core.BlobReference("optimizer_iteration") in checkpoint_params)
  
      def test_net_conversion_and_append_net(self):
          other = model_helper.ModelHelper()
          fc1 = brew.fc(other, "data", "other_fc1", dim_in=3*227*227, dim_out=10)
          fc2 = brew.fc(other, fc1, "other_fc2", dim_in=10, dim_out=10)
          brew.fc(other, fc2, "other_fc3", dim_in=10, dim_out=10)
  
          def add_input_ops(model):
              model.net.UniformFill([], ["data"], shape=[4, 227, 227, 3])
              model.net.UniformFill([], ["label"], shape=[4])
  
          def add_model_ops(model, loss_scale):
              model.NHWC2NCHW("data", "data_nchw")
              model.Conv("data_nchw", 'conv1', 3, 64,
                         weight_init=("MSRAFill", {}), kernel=7,
                         stride=2, pad=3, no_bias=0)
              model.SpatialBN('conv1', 'conv1_spatbn_relu', 64, epsilon=1e-3, is_test=False)
              model.Relu('conv1_spatbn_relu', 'conv1_spatbn_relu')
              model.MaxPool('conv1_spatbn_relu', 'pool1', kernel=3, stride=2)
              model.FC('pool1', 'fc', dim_in=(64 * 56 * 56), dim_out=10)
  
              # Append the net and param_init_net of the other model
              appendnet = data_parallel_model.ConvertNetForDevice(other.net)
              model.net.AppendNet(appendnet)
  
              model.param_init_net.AppendNet(
                  data_parallel_model.ConvertNetForDevice(other.param_init_net))
  
              model.Sigmoid('fc', 'fc_sigm')
              model.Softmax('fc_sigm', 'softmax')
              loss = model.AveragedLoss('softmax', 'loss')
              return [loss]
  
          def add_optimizer(model):
              optimizer.build_sgd(model, 0.1, policy="fixed", momentum=0.9)
  
          model = cnn.CNNModelHelper(
              order="NCHW",
              name="test",
          )
          data_parallel_model.Parallelize_CPU(
              model,
              input_builder_fun=add_input_ops,
              forward_pass_builder_fun=add_model_ops,
              optimizer_builder_fun=add_optimizer,
              devices=range(4)
          )
  
          # Just create and run net and confirm no exception is thrown
          workspace.RunNetOnce(model.param_init_net)
          workspace.CreateNet(model.net)
          workspace.RunNet(model.net)
  
      def test_synchronization_barrier(self):
          def run(comm_rank, comm_size, tmpdir):
              def add_input_ops(model):
                  pass
  
              def add_model_ops(model, loss_scale):
                  return []
  
              def add_optimizer(model):
                  pass
  
              store_handler = "store_handler"
              workspace.RunOperatorOnce(
                  core.CreateOperator(
                      "FileStoreHandlerCreate",
                      [],
                      [store_handler],
                      path=tmpdir))
              rendezvous = dict(
                  kv_handler=store_handler,
                  shard_id=comm_rank,
                  num_shards=comm_size,
                  engine='GLOO',
              )
  
              model = cnn.CNNModelHelper(
                  order="NHWC",
                  name="test",
              )
              data_parallel_model.Parallelize_CPU(
                  model,
                  input_builder_fun=add_input_ops,
                  forward_pass_builder_fun=add_model_ops,
                  optimizer_builder_fun=add_optimizer,
                  devices=[1, 2, 3],
                  rendezvous=rendezvous
              )
              data_parallel_model.RunInitNet(model)
  
              for _ in range(2):
                  data_parallel_model.Synchronize(model)
  
          with TemporaryDirectory() as tmpdir:
              self.run_test_locally(
                  run,
                  comm_size=2,
                  device_option=None,
                  tmpdir=tmpdir)
  
      def test_pre_train_synchronization_barrier(self):
          def run(comm_rank, comm_size, tmpdir):
              def add_input_ops(model):
                  pass
  
              def add_model_ops(model, loss_scale):
                  return []
  
              def add_optimizer(model):
                  pass
  
              workspace.ResetWorkspace()
              store_handler = "store_handler"
              workspace.RunOperatorOnce(
                  core.CreateOperator(
                      "FileStoreHandlerCreate",
                      [],
                      [store_handler],
                      path=tmpdir))
              rendezvous = dict(
                  kv_handler=store_handler,
                  shard_id=comm_rank,
                  num_shards=comm_size,
                  engine='GLOO',
              )
  
              model = cnn.CNNModelHelper(
                  order="NHWC",
                  name="test",
              )
              # Set network timeout to 2 seconds, and add a 3 seconds
              # sleep for 1 host.  Make sure there is no timeout on the
              # second RunNet.
              data_parallel_model._DEFAULT_TIMEOUT_SEC = 2
              data_parallel_model.Parallelize_CPU(
                  model,
                  input_builder_fun=add_input_ops,
                  forward_pass_builder_fun=add_model_ops,
                  optimizer_builder_fun=add_optimizer,
                  devices=[1, 2, 3],
                  rendezvous=rendezvous,
                  barrier_net_timeout_sec=5
              )
              data_parallel_model.RunInitNet(model)
              data_parallel_model.RunNet(model, 2)
              if comm_rank == 0:
                  time.sleep(data_parallel_model._DEFAULT_TIMEOUT_SEC)
              data_parallel_model.RunNet(model, 2)
  
          with TemporaryDirectory() as tmpdir:
              self.run_test_locally(
                  run,
                  comm_size=2,
                  device_option=None,
                  tmpdir=tmpdir)
  
      def test_device_scope_check(self):
          with self.assertRaises(AssertionError):
              with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, 0)):
                  data_parallel_model.Parallelize_GPU(None, None, None)
  
      def test_net_transformer_function(self):
          devices = [1, 2, 3]
  
          def add_input_ops(model):
              model.param_init_net.UniformFill([], ["data"], shape=[32, 8])
  
          def add_optimizer(model):
              optimizer.build_sgd(model, 0.1)
  
          def add_model_ops(model, loss_scale):
              fc1 = brew.fc(model, "data", "fc1", dim_in=8, dim_out=8)
              return [fc1]
  
          kwargs = {
              'input_builder_fun': add_input_ops,
              'forward_pass_builder_fun': add_model_ops,
              'devices': devices,
          }
  
          # assert that the transformer is called for both train and test cases
          transform = Mock()
          kwargs['net_transformer_fun'] = transform
          model = model_helper.ModelHelper(name="r", init_params=False)
          data_parallel_model.Parallelize_CPU(model, **kwargs)
          self.assertTrue(transform.called)
          self.assertEqual(transform.call_count, 1)
  
          transform = Mock()
          kwargs['net_transformer_fun'] = transform
          kwargs['optimizer_builder_fun'] = add_optimizer
          model = model_helper.ModelHelper(name="r", init_params=True)
          data_parallel_model.Parallelize_CPU(model, **kwargs)
          self.assertTrue(transform.called)
          self.assertEqual(transform.call_count, 1)
  
      @given(seed=st.integers(0, 65535), batch_size=st.integers(1, 20))
      def test_multi_device_bn_op_level_cpu(self, seed, batch_size):
          self._bn_check_op_level("cpu", seed, batch_size)
  
      @unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")
      @unittest.skipIf(workspace.NumCudaDevices() < 2, "Need at least 2 GPUs.")
      @given(seed=st.integers(0, 65535), batch_size=st.integers(1, 20))
      def test_multi_device_bn_op_level_gpu(self, seed, batch_size):
          self._bn_check_op_level("gpu", seed, batch_size)
  
      def _bn_check_op_level(self, device_type, seed, batch_size):
          '''
          Test multi device batch normalization at the operation level. This is
          done by checking the outputs of batch normalization and its gradient
          operator. We compare values produced with our manually calculated
          batch normalization values and gradients.
          '''
          devices = [0, 1]
          epsilon = 1e-3
          tolerance = 1e-3
  
          def _test_forward_pass(x, devices, device_type, scale, bias, epsilon):
              x_concat = np.concatenate(x)
              mean = np.mean(x_concat, axis=0)
              var = np.var(x_concat, axis=0)
              for device in devices:
                  x_i = x[device]
                  x_hat = (x_i - mean) / (np.sqrt(var + epsilon))
                  expected_out = scale * x_hat + bias
                  spatial_out = workspace.FetchBlob(
                      "{}_{}/bn_out".format(device_type, device))
                  rel_error = np.linalg.norm(spatial_out - expected_out) \
                              / np.linalg.norm(expected_out)
                  self.assertTrue(rel_error < 0.005)
  
          def _test_backward_pass(x, devices, device_type, scale, tolerance):
              dBias_arr = []
              dY_arr = []
              dGamma_arr = []
              num_devices = len(devices)
              mean = np.array(workspace.FetchBlob(
                  "{}_0/bn_out_sm".format(device_type)), dtype=np.float32)
              inv_var = np.array(workspace.FetchBlob(
                  "{}_0/bn_out_siv".format(device_type)), dtype=np.float32)
  
              # dBias
              # Sum dBias values over all devices to find the average gradient
              for device in devices:
                  dY_blob = workspace.FetchBlob(
                      "{}_{}/bn_out_grad".format(device_type, device))
                  dY = np.array(dY_blob, dtype=np.float32)
                  dY_arr.append(dY)
                  dBias_arr.append(np.array(np.sum(dY, axis=0), dtype=np.float32))
              dBias = np.sum(dBias_arr, dtype=np.float32)
              dBias_avg = dBias / num_devices
              for device in devices:
                  dBiasActual = np.sum(workspace.FetchBlob("{}_{}/bn_out_b_grad"
                      .format(device_type, device)), dtype=np.float32)
                  self.assertTrue(np.isclose([dBiasActual], [dBias], atol=tolerance))
  
              # dGamma
              # Sum dGamma values over all devices to find the average gradient
              for device in devices:
                  dGamma = np.sum((x[device] - mean) * inv_var * dY_arr[device],
                      axis=0, dtype=np.float32)
                  dGamma_arr.append(dGamma)
              dGamma = np.sum(dGamma_arr, axis=0, dtype=np.float32)
              dGamma_avg = dGamma / num_devices
              for device in devices:
                  dGammaActual = workspace.FetchBlob(
                      "{}_{}/bn_out_s_grad".format(device_type, device))
                  self.assertTrue(np.isclose([dGamma], [dGammaActual], atol=tolerance))
  
              # dX
              scale_inv_var = scale * inv_var / batch_size
              for device in devices:
                  dX = scale_inv_var * (dY_arr[device] * batch_size - dBias_avg
                      - (x[device] - mean) * dGamma_avg * inv_var)
                  dX_actual = workspace.FetchBlob(
                      "{}_{}/tanh_grad".format(device_type, device))
                  self.assertTrue(np.isclose([dX], [dX_actual], atol=tolerance).all())
  
          def add_input_ops(model):
              for device in devices:
                  data = np.random.rand(batch_size, 1, 1, 1).astype(np.float32)
                  workspace.FeedBlob("{}_{}/data".format(device_type, device), data)
  
          def add_model_ops(model, loss_scale):
              if device_type == "gpu":
                  model.CopyCPUToGPU("data", "device_data")
                  model.Tanh("device_data", "tanh")
              else:
                  model.Tanh("data", "tanh")
              model.SpatialBN("tanh", "bn_out", 1, epsilon=epsilon, is_test=False)
              model.Sqr("bn_out", "sqr")
              loss = model.SumElements("sqr", "loss")
              return [loss]
  
          def add_optimizer(model):
              return optimizer.build_sgd(model, 0.1)
  
          np.random.seed(seed)
          workspace.ResetWorkspace()
          model = cnn.CNNModelHelper(
              order="NCHW",
              name="test"
          )
          data_parallel_model.Parallelize(
              model,
              input_builder_fun=add_input_ops,
              forward_pass_builder_fun=add_model_ops,
              optimizer_builder_fun=add_optimizer,
              devices=devices,
              cpu_device=device_type == "cpu",
              shared_model=False,
              combine_spatial_bn=True,
          )
  
          workspace.RunNetOnce(model.param_init_net)
          scale = workspace.FetchBlob("{}_0/bn_out_s".format(device_type))
          bias = workspace.FetchBlob("{}_0/bn_out_b".format(device_type))
          workspace.RunNetOnce(model.net)
  
          x = []
          for device in devices:
              x_blob = workspace.FetchBlob("{}_{}/tanh".format(device_type, device))
              x_i = np.array(x_blob, dtype=np.float32)
              x.append(x_i)
  
          _test_forward_pass(x, devices, device_type, scale, bias, epsilon)
          _test_backward_pass(x, devices, device_type, scale, tolerance)
  
      @given(seed=st.integers(0, 65535), batch_size=st.integers(1, 20))
      def test_multi_device_bn_net_lvl_cpu(self, seed, batch_size):
          if batch_size % 2 == 1:
              batch_size += 1
          self._test_multi_device_bn_net_lvl("cpu", seed, batch_size)
  
      @unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")
      @unittest.skipIf(workspace.NumCudaDevices() < 2, "Need at least 2 GPUs.")
      @given(seed=st.integers(0, 65535), batch_size=st.integers(1, 20))
      def test_multi_device_bn_net_lvl_gpu(self, seed, batch_size):
          if batch_size % 2 == 1:
              batch_size += 1
          self._test_multi_device_bn_net_lvl("gpu", seed, batch_size)
  
      def _test_multi_device_bn_net_lvl(self, device_type, seed, batch_size):
          '''
          Test multi device batch normalization at the net level. This is done
          by verifying that the final batch normalization outputs and the
          gradient outputs from multiple devices are the same as those produced
          from a single device
          '''
  
          # Verify that the gradients calculated over multiple devices are the
          # same as the gradients calculated over one device. These values should
          # be equivalent because combine_spatial_bn sums values over all devices
          def _verify_bn_outputs(
              devices,
              device_type,
              tolerance,
              single_device_bn_out,
              two_device_bn_out_vals,
              single_device_grads,
              two_device_grads,
          ):
              two_device_bn_out = np.concatenate(two_device_bn_out_vals)
              self.assertTrue(np.isclose(
                  [single_device_bn_out], [two_device_bn_out], atol=tolerance).all())
  
              # Scalar and Bias gradients should be the same across devices
              gradient_names = ["bn_out_s_grad", "bn_out_b_grad"]
              for name in gradient_names:
                  expected_grad = single_device_grads[name]
                  for device in devices:
                      actual_grad = two_device_grads[device][name]
                      self.assertTrue(
                          np.isclose([actual_grad], [expected_grad], atol=tolerance))
  
              # Expected tanh_grad should be the combined tanh_grad vectors
              # across the devices
              first_grad = two_device_grads[0]["tanh_grad"]
              second_grad = two_device_grads[1]["tanh_grad"]
              actual_grad = np.concatenate([first_grad, second_grad])
              expected_grad = single_device_grads["tanh_grad"]
              rel_error = np.linalg.norm(actual_grad - expected_grad) \
                  / np.linalg.norm(expected_grad)
              self.assertTrue(rel_error < 1e-3)
  
          def _create_model(multiple_devices):
              def add_input_ops_no_combine(model):
                  workspace.FeedBlob("{}_0/data".format(device_type), data)
  
              def add_input_ops_combine(model):
                  half = int(batch_size / 2)
                  workspace.FeedBlob("{}_0/data".format(device_type), data[:half])
                  workspace.FeedBlob("{}_1/data".format(device_type), data[half:])
  
              def add_model_ops(model, loss_scale):
                  if device_type == "gpu":
                      model.CopyCPUToGPU("data", "device_data")
                      model.Tanh("device_data", "tanh")
                  else:
                      model.Tanh("data", "tanh")
                  model.SpatialBN("tanh", "bn_out", 1, epsilon=epsilon, is_test=False)
                  model.Sqr("bn_out", "sqr")
                  loss = model.SumElements("sqr", "loss")
                  return [loss]
  
              def add_optimizer(model):
                  return optimizer.build_sgd(model, 0.1)
  
              if multiple_devices:
                  input_fun = add_input_ops_combine
                  devices = [0, 1]
                  combine_spatial_bn = True
              else:
                  input_fun = add_input_ops_no_combine
                  devices = [0]
                  combine_spatial_bn = False
              model = cnn.CNNModelHelper(
                  order="NCHW",
                  name="test"
              )
              data_parallel_model.Parallelize(
                  model,
                  input_builder_fun=input_fun,
                  forward_pass_builder_fun=add_model_ops,
                  optimizer_builder_fun=add_optimizer,
                  devices=devices,
                  cpu_device=device_type == "cpu",
                  shared_model=False,
                  combine_spatial_bn=combine_spatial_bn,
              )
              return model
  
          devices = [0, 1]
          epsilon = 1e-3
          tolerance = 1e-3
          # We are generating random data
          np.random.seed(seed)
          data = np.random.rand(batch_size, 1, 1, 1).astype(np.float32)
          data = np.reshape(data, (batch_size, 1, 1, 1))
  
          # Get values calculated without combine_spatial_bn
          workspace.ResetWorkspace()
          model_no_combine = _create_model(multiple_devices=False)
          workspace.RunNetOnce(model_no_combine.param_init_net)
          workspace.RunNetOnce(model_no_combine.net)
          single_device_bn_out = workspace.FetchBlob("{}_0/bn_out".format(device_type))
          single_device_grads = {}
          single_device_grads["bn_out_s_grad"] = workspace.FetchBlob(
              "{}_0/bn_out_s_grad".format(device_type))
          single_device_grads["bn_out_b_grad"] = workspace.FetchBlob(
              "{}_0/bn_out_b_grad".format(device_type))
          single_device_grads["tanh_grad"] = workspace.FetchBlob(
              "{}_0/tanh_grad".format(device_type))
  
          # Get values calculated over multiple devices with combine_spatial_bn true
          workspace.ResetWorkspace()
          model_combine = _create_model(multiple_devices=True)
          workspace.RunNetOnce(model_combine.param_init_net)
          workspace.RunNetOnce(model_combine.net)
          two_device_bn_out_vals = []
          two_device_grads = {}
          for device in devices:
              bn_out_blob = "{}_{}/bn_out".format(device_type, device)
              two_device_bn_out_vals.append(workspace.FetchBlob(bn_out_blob))
              two_device_grads[device] = {}
              two_device_grads[device]["bn_out_s_grad"] = workspace.FetchBlob(
                  "{}_{}/bn_out_s_grad".format(device_type, device))
              two_device_grads[device]["bn_out_b_grad"] = workspace.FetchBlob(
                  "{}_{}/bn_out_b_grad".format(device_type, device))
              two_device_grads[device]["tanh_grad"] = workspace.FetchBlob(
                  "{}_{}/tanh_grad".format(device_type, device))
  
          # Check to see if the combined values are equivalent
          _verify_bn_outputs(
              devices,
              device_type,
              tolerance,
              single_device_bn_out,
              two_device_bn_out_vals,
              single_device_grads,
              two_device_grads
          )
  
  class RecurrentNetworkParallelTest(TestCase):
  
      def run_model(self, devices, gpu):
  
          '''
          Helper function for test_equiv
          '''
          def input_builder_fun(model):
              return None
  
          def model_build_fun(model, loss_scale):
              workspace.FeedBlob(
                  core.ScopedBlobReference("seq_lengths"),
                  np.array([self.T] * self.batch_per_device, dtype=np.int32)
              )
              model.param_init_net.ConstantFill(
                  [],
                  "hidden_init",
                  value=0.0,
                  shape=[1, self.batch_per_device, self.hidden_dim]
              )
              model.param_init_net.ConstantFill(
                  [],
                  "cell_init",
                  value=0.0,
                  shape=[1, self.batch_per_device, self.hidden_dim]
              )
  
              output, _last_hidden, _, _last_state, = rnn_cell.LSTM(
                  model=model,
                  input_blob="data",
                  seq_lengths="seq_lengths",
                  initial_states=("hidden_init", "cell_init"),
                  dim_in=self.input_dim,
                  dim_out=self.hidden_dim,
                  scope="partest",
              )
  
              # A silly loss function
              loss = model.AveragedLoss(
                  model.Sub([output, "target"], "dist"),
                  "loss",
              )
              loss = model.Scale(loss, "loss_scaled", scale=loss_scale)
              return [loss]
  
          def param_update_fun(model):
              ITER = model.Iter("ITER")
              LR = model.net.LearningRate(
                  [ITER],
                  "LR",
                  base_lr=(-0.1),
                  policy="fixed",
              )
              ONE = model.param_init_net.ConstantFill(
                  [], "ONE", shape=[1], value=1.0,
              )
              for param in model.GetParams():
                  param_grad = model.param_to_grad[param]
                  model.WeightedSum([param, ONE, param_grad, LR], param)
  
              assert len(model.GetParams()) == len(model.params) // len(model._devices)
  
          workspace.ResetWorkspace()
          model = cnn.CNNModelHelper(
              name="recurrent_test{}".format(devices),
          )
  
          self.T = 8
          self.batch_size = 64
          self.input_dim = 8
          self.hidden_dim = 31
          self.batch_per_device = self.batch_size // len(devices)
  
          data_parallel_model.Parallelize(
              model,
              input_builder_fun=input_builder_fun,
              forward_pass_builder_fun=model_build_fun,
              param_update_builder_fun=param_update_fun,
              devices=devices,
              optimize_gradient_memory=True,
              cpu_device=not gpu,
          )
  
          # Change all initialization to be ConstantFills so that
          # the everything is deterministic
          for op in model.param_init_net.Proto().op:
              if op.type.endswith('Fill'):
                  op.type = 'ConstantFill'
  
          # Each run has same input, independent of number of gpus
          np.random.seed(20150210)
          for i in range(0, 10):
              full_data = np.random.rand(self.T, self.batch_size, self.input_dim)
              full_target = np.random.rand(
                  self.T, self.batch_size, self.hidden_dim
              )
  
              for (j, g) in enumerate(devices):
                  st = j * self.batch_per_device
                  en = st + self.batch_per_device
                  data = full_data[:, st:en, :].astype(np.float32)
                  targets = full_target[:, st:en, :].astype(np.float32)
                  with core.DeviceScope(core.DeviceOption(model._device_type, g)):
                      workspace.FeedBlob(
                          "{}_{}/data".format(model._device_prefix, g), data
                      )
                      workspace.FeedBlob(
                          "{}_{}/target".format(model._device_prefix, g), targets
                      )
  
              if i == 0:
                  workspace.RunNetOnce(model.param_init_net)
                  workspace.CreateNet(model.net)
  
              workspace.RunNet(model.net.Proto().name)
  
          return workspace.FetchBlob("{}_0/partest/i2h_w".format(model._device_prefix))
  
      @unittest.skip("Test is flaky: https://github.com/pytorch/pytorch/issues/10322")
      def test_equiv_recurrent(self):
          '''
          Test that the model produces exactly same results given
          total batchsize, independent of number of GPUs/CPUs.
          '''
          for gpu in [True, False]:
              if gpu and not workspace.has_gpu_support:
                  continue
              result_2gpus = self.run_model([0, 1], gpu)
              result_1gpus = self.run_model([0], gpu)
  
              self.assertTrue(np.allclose(result_1gpus, result_2gpus))
  
              if not gpu or workspace.NumCudaDevices() >= 4:
                  result_4gpus = self.run_model(list(range(4)), gpu)
                  self.assertTrue(np.allclose(result_1gpus, result_4gpus))
  
              if not gpu or workspace.NumCudaDevices() >= 8:
                  result_8gpus = self.run_model(list(range(8)), gpu)
                  self.assertTrue(np.allclose(result_1gpus, result_8gpus))
  
  
  @unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")
  @unittest.skipIf(workspace.NumCudaDevices() < 2, "Need at least 2 GPUs.")
  class SparseDataParallelModelTest(TestCase):
  
      '''
      Create and run the model. We try with both storing indices for gather
      on CPU and on GPU
      '''
      def run_model(self, V, gpu_devices, cpu_indices):
  
          def input_builder_fun(model):
              return None
  
          def model_build_fun(model, loss_scale):
              if cpu_indices:
                  with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
                      gathered_cpu = model.net.Gather(
                          [self.vecs, 'indices'], 'gathered_cpu')
  
                  gathered = model.CopyCPUToGPU(gathered_cpu, "gathered")
              else:
                  gpu_vecs = model.param_init_net.CopyCPUToGPU(
                      self.vecs, "gpuvecs",
                  )
                  model.params.append(gpu_vecs)
                  gathered = model.net.Gather([gpu_vecs, 'indices'], 'gathered')
              flattened = model.Flatten(gathered, "flattened")
              fc = model.FC(flattened, "fc", 16 * 16, 1,
                            ("ConstantFill", {}), ("ConstantFill", {}))
              fc_fl = model.FlattenToVec(fc, "fc_fl")
              sigm = model.Sigmoid(fc_fl, "sigm")
              sq = model.SquaredL2Distance([sigm, "label"], "sq")
              loss = model.AveragedLoss(sq, "loss")
              loss = model.Scale(loss, scale=loss_scale)
              return [loss]
  
          def param_update_fun(model):
              ONE = model.param_init_net.ConstantFill(
                  [], "ONE", shape=[1], value=1.0,
              )
              LR = model.CopyCPUToGPU(self.LR, "LR")
              for param in model.GetParams():
                  param_grad = model.param_to_grad[param]
                  if not isinstance(param_grad, core.GradientSlice):
                      model.WeightedSum([param, ONE, param_grad, LR], param)
                  else:
                      param_momentum = model.param_init_net.ConstantFill(
                          [param],
                          param + '_momentum',
                          value=0.0,
                      )
                      model.net.SparseMomentumSGDUpdate(
                          [
                              param_grad.values,
                              param_momentum,
                              LR,
                              param,
                              param_grad.indices,
                          ],
                          [
                              param_grad.values, param_momentum, param
                          ],
                          momentum=0.1,
                          nesterov=0,
                      )
  
          workspace.ResetWorkspace()
          model = cnn.CNNModelHelper(
              order="NHWC",
              name="sparse_test{}".format(gpu_devices),
          )
  
          with core.NameScope("cpu"):
              with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
                  self.ITER = model.Iter("ITER")
                  self.LR = model.net.LearningRate(
                      [self.ITER],
                      "LR",
                      base_lr=(-0.1),
                      policy="fixed",
                  )
                  self.vecs = model.param_init_net.UniformFill(
                      [], "vecs", shape=[V, 16])
                  if cpu_indices:
                      model.params.append(self.vecs)
                  self.ONE_CPU = model.param_init_net.ConstantFill(
                      [], "ONE_CPU", shape=[1], value=1.0,
                  )
  
          data_parallel_model.Parallelize_GPU(
              model,
              input_builder_fun=input_builder_fun,
              forward_pass_builder_fun=model_build_fun,
              param_update_builder_fun=param_update_fun,
              devices=gpu_devices,
          )
  
          # Update the vecs
          if cpu_indices:
              with core.NameScope("cpu"):
                  with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
                      for param in model.GetParams():
                          param_grad = model.param_to_grad[param]
                          model.ScatterWeightedSum([param, self.ONE_CPU,
                                                    param_grad.indices,
                                                    param_grad.values,
                                                    self.LR],
                                                    self.vecs)
          else:
              with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, 0)):
                  model.CopyGPUToCPU("gpu_0/gpuvecs", self.vecs)
  
          np.random.seed(2603)
  
          # Each run has same input, independent of number of gpus
          batch_size = 64
          for i in range(0, 10):
              full_indices = np.random.permutation(V)[:batch_size * 16].reshape(
                  batch_size, 16
              )
              full_labels = full_indices[:, 0] % 2
              batch_per_device = batch_size // len(gpu_devices)
  
              for (j, g) in enumerate(gpu_devices):
                  st = j * batch_per_device
                  en = st + batch_per_device
                  indices = full_indices[st:en, :].astype(np.int32)
                  labels = full_labels[st:en].astype(np.float32)
  
                  device_for_indices = core.DeviceOption(caffe2_pb2.CPU)
                  if not cpu_indices:
                      device_for_indices = core.DeviceOption(workspace.GpuDeviceType, g)
  
                  with core.DeviceScope(device_for_indices):
                      workspace.FeedBlob("gpu_{}/indices".format(g), indices)
  
                  with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, g)):
                      workspace.FeedBlob("gpu_{}/label".format(g), labels)
  
              if i == 0:
                  workspace.RunNetOnce(model.param_init_net)
                  # Force vecs to be same on all runs
                  orig_vecs = np.random.rand(V, 16).astype(np.float32)
                  workspace.FeedBlob(
                      self.vecs,
                      orig_vecs
                  )
                  if not cpu_indices:
                      for g in gpu_devices:
                          workspace.FeedBlob(
                              "gpu_{}/gpuvecs".format(g),
                              orig_vecs,
                              device_option=core.DeviceOption(workspace.GpuDeviceType, g),
                          )
                  workspace.CreateNet(model.net)
  
              workspace.RunNet(model.net.Proto().name)
              if len(gpu_devices) == 2:
                  if not cpu_indices:
                      idx = workspace.FetchBlob("gpu_0/indices")
                      idx = list(idx.flatten())
                      n = len(idx)
                      nu = len(set(idx))
                      assert n == nu, "We cannot have duplicate indices"
  
          # Sanity check to see the vecs were updated
          self.assertFalse(
              np.allclose(workspace.FetchBlob(self.vecs), orig_vecs))
          return [workspace.FetchBlob(self.vecs if cpu_indices else "gpu_0/gpuvecs"),
                  workspace.FetchBlob("gpu_0/fc_w")]
  
      def _test_equiv_sparse(self, cpu_indices):
          '''
              Test that the model produces exactly same results given
              total batchsize, independent of number of GPUs.
          '''
          V = 10000
          result_2gpus = self.run_model(V, [0, 1], cpu_indices)
          result_1gpus = self.run_model(V, [0], cpu_indices)
  
          self.assertTrue(np.allclose(result_1gpus[0], result_2gpus[0]))
          self.assertTrue(np.allclose(result_1gpus[1], result_2gpus[1]))
  
          if workspace.NumCudaDevices() >= 4:
              result_4gpus = self.run_model(V, list(range(4)), cpu_indices)
              self.assertTrue(np.allclose(result_1gpus[0], result_4gpus[0]))
              self.assertTrue(np.allclose(result_1gpus[1], result_4gpus[1]))
  
          if workspace.NumCudaDevices() >= 8:
              result_8gpus = self.run_model(V, list(range(8)), cpu_indices)
              self.assertTrue(np.allclose(result_1gpus[0], result_8gpus[0]))
              self.assertTrue(np.allclose(result_1gpus[1], result_8gpus[1]))
  
      def test_equiv_sparse(self):
          self._test_equiv_sparse(True)
          self._test_equiv_sparse(False)
  
  
  @unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")
  @unittest.skipIf(workspace.NumGpuDevices() < 2, "Need at least 2 GPUs.")
  class ParallelizeBMUFTest(TestCase):
  
      def _run_model(self, gpu_devices):
          '''
          Helper function for test_equiv
          '''
          def input_builder_fun(model):
              return None
  
      def _model_build_fun(self, model, loss_scale):
          fc = model.FC(
              "data", "fc", 16, 1, ("ConstantFill", {}), ("ConstantFill", {})
          )
          fc_fl = model.FlattenToVec(fc, "fc_fl")
          sigm = model.Sigmoid(fc_fl, "sigm")
          sq = model.SquaredL2Distance([sigm, "label"], "sq")
          loss = model.AveragedLoss(sq, "loss")
          loss = model.Scale(loss, scale=loss_scale)
  
          return [loss]
  
      def _param_update_fun(self, model):
          ITER = model.Iter("ITER")
          LR = model.net.LearningRate(
              [ITER],
              "LR",
              base_lr=(-0.1),
              policy="fixed",
          )
          ONE = model.param_init_net.ConstantFill(
              [], "ONE", shape=[1], value=1.0,
          )
          for param in model.GetParams():
              grad = model.param_to_grad[param]
              model.WeightedSum([param, ONE, grad, LR], param)
  
      def _generate_data(self, devices, device_type, device_prefix):
          np.random.seed(26)
          # Each run has same input, independent of number of gpus
          batch_size = 64
          for _ in range(0, 10):
              full_data = np.random.rand(batch_size, 16)
              full_labels = np.round(full_data[:, 0])
              batch_per_device = batch_size // len(devices)
  
              for (j, g) in enumerate(devices):
                  st = j * batch_per_device
                  en = st + batch_per_device
                  data = full_data[st:en, :].astype(np.float32)
                  labels = full_labels[st:en].astype(np.float32)
                  with core.DeviceScope(core.DeviceOption(device_type, g)):
                      workspace.FeedBlob("{}_{}/data".format(device_prefix, g), data)
                      workspace.FeedBlob("{}_{}/label".format(device_prefix, g), labels)
  
      @given(
          cpu_device=st.booleans()
      )
      def test_parallelize_bmuf(self, cpu_device):
          assume(cpu_device or workspace.has_gpu_support or workspace.has_hip_support)
  
          workspace.ResetWorkspace()
  
          model = cnn.CNNModelHelper(
              order="NHWC",
              name="test"
          )
          devices = [0, 1]
  
          def input_builder_fun(model):
              return None
  
          if not cpu_device:
              device_type = workspace.GpuDeviceType
              device_prefix = "gpu"
          else:
              device_type = caffe2_pb2.CPU
              device_prefix = "cpu"
          self._generate_data(devices, device_type, device_prefix)
  
          data_parallel_model.Parallelize_BMUF(
              model,
              input_builder_fun,
              self._model_build_fun,
              self._param_update_fun,
              devices=devices,
              cpu_device=cpu_device
          )
  
          data_parallel_model.RunInitNet(model)
  
          # Check initial momentum params are zeros
          self.assertEqual(
              list(viewkeys(model._device_grouped_blobs)), ['fc_w', 'fc_b']
          )
          self.assertEqual(workspace.FetchBlob('{}_0/fc_b_v'.format(device_prefix)), 0)
          np.testing.assert_equal(
              workspace.FetchBlob('{}_0/fc_w_v'.format(device_prefix)),
              np.zeros(16).astype(np.float32).reshape(1, 16)
          )
  
          # Run the algorithm for one iteration to have non-zero params.
          data_parallel_model.RunNet(model, 1)
  
          # Save iteration momentum and post local update params
          v_b_ = workspace.FetchBlob('{}_0/fc_b_v'.format(device_prefix))
          v_w_ = workspace.FetchBlob('{}_0/fc_w_v'.format(device_prefix))
  
          workspace.RunNetOnce(model.net)
  
          b_0_ = workspace.FetchBlob('{}_0/fc_b'.format(device_prefix))
          w_0_ = workspace.FetchBlob('{}_0/fc_w'.format(device_prefix))
          b_1_ = workspace.FetchBlob('{}_1/fc_b'.format(device_prefix))
          w_1_ = workspace.FetchBlob('{}_1/fc_w'.format(device_prefix))
  
          # Compute block gradients.
          b_g_ = workspace.FetchBlob('{}_0/fc_b_g'.format(device_prefix))
          w_g_ = workspace.FetchBlob('{}_0/fc_w_g'.format(device_prefix))
          workspace.RunNetOnce(model._global_model_param_updates_net)
  
          g_b = (b_0_ + b_1_) / 2 - b_g_
          g_w = (w_0_ + w_1_) / 2 - w_g_
          v_b = workspace.FetchBlob('{}_0/fc_b_v'.format(device_prefix))
          v_w = workspace.FetchBlob('{}_0/fc_w_v'.format(device_prefix))
  
          w_g = workspace.FetchBlob('{}_0/fc_w_g'.format(device_prefix))
          b_g = workspace.FetchBlob('{}_0/fc_b_g'.format(device_prefix))
          w_0 = workspace.FetchBlob('{}_0/fc_w'.format(device_prefix))
          b_0 = workspace.FetchBlob('{}_0/fc_b'.format(device_prefix))
          w_1 = workspace.FetchBlob('{}_1/fc_w'.format(device_prefix))
          b_1 = workspace.FetchBlob('{}_1/fc_b'.format(device_prefix))
  
          # Check momentum update step
          np.testing.assert_equal(v_b, 0.5 * v_b_ + g_b)
          np.testing.assert_equal(v_w, 0.5 * v_w_ + g_w)
  
          np.testing.assert_equal(w_g, w_0)
          np.testing.assert_equal(w_g, w_1)
          np.testing.assert_equal(b_g, b_0)
          np.testing.assert_equal(b_g, b_1)
  
          # Check params update step
          np.testing.assert_equal(w_0, w_g_ + v_w)
          np.testing.assert_equal(b_0, b_g_ + v_b)
  
  
  @unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")
  @unittest.skipIf(workspace.NumGpuDevices() < 2, "Need at least 2 GPUs.")
  class SparseDataParallelModelTestWithSharedIndices(TestCase):
  
      '''
      Create and run the model. We try with both storing indices for gather
      on CPU and on GPU
      '''
      def run_model(self, V, gpu_devices):
  
          def input_builder_fun(model):
              return None
  
          def model_build_fun(model, loss_scale):
              gpu_vecs_gathered = []
              gpu_vecs = []
              for num, vec in enumerate(self.vecs):
                  gpu_vec = model.param_init_net.CopyCPUToGPU(
                      vec, 'gpuvec_{}'.format(num),
                  )
                  if num != 2:
                      model.params.append(gpu_vec)
                  gpu_vecs.append(gpu_vec)
              for num, gpu_vec in enumerate(gpu_vecs):
                  gpu_vec_gathered = model.net.Gather(
                      [gpu_vec, 'indices'],
                      ['gpu_vec_gathered_{}'.format(num)]
                  )
                  gpu_vecs_gathered.append(gpu_vec_gathered)
  
              assert len(gpu_vecs_gathered) == 3
  
              fc = model.net.FC(
                  [
                      gpu_vecs_gathered[2],
                      gpu_vecs_gathered[0],
                      gpu_vecs_gathered[1],
                  ],
                  ['fc'],
              )
              _, loss = model.net.SoftmaxWithLoss(
                  [fc, 'label'],
                  ['ce_loss', 'avg_loss'],
                  only_loss=True,
              )
              loss = model.Scale(loss, scale=loss_scale)
              model.net.Print(loss, [], limit=10)
              return [loss]
  
          def param_update_fun(model):
              ONE = model.param_init_net.ConstantFill(
                  [], "ONE", shape=[1], value=1.0,
              )
              LR = model.CopyCPUToGPU(self.LR, "LR")
              for param in model.GetParams():
                  param_grad = model.param_to_grad[param]
                  if not isinstance(param_grad, core.GradientSlice):
                      model.WeightedSum([param, ONE, param_grad, LR], param)
                  else:
                      model.net.ScatterWeightedSum(
                          [
                              param,
                              ONE,
                              param_grad.indices,
                              param_grad.values,
                              ONE,
                          ],
                          param,
                      )
  
          workspace.ResetWorkspace()
          model = cnn.CNNModelHelper(
              order="NHWC",
              name="sparse_test{}".format(gpu_devices),
          )
          batch_size = 32
          batch_per_device = batch_size // len(gpu_devices)
  
          with core.NameScope("cpu"):
              with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
                  self.ITER = model.Iter("ITER")
                  self.LR = model.net.LearningRate(
                      [self.ITER],
                      "LR",
                      base_lr=(-0.1),
                      policy="fixed",
                  )
                  '''
                  self.vecs consists of 3 big blobs on which we call Gather:
                  1) FC weights, shape=(V, 16)
                  2) FC bias, shape=(V)
                  3) FC input, shape=(batch_per_device, 16)
                  '''
                  self.vecs = [
                      model.param_init_net.UniformFill(
                          [], "vec_{}".format(num), shape=[V, 16])
                      for num in range(2)
                  ]
                  self.vecs.append(
                      model.param_init_net.UniformFill(
                          [],
                          "vec_2", shape=[batch_per_device, 16]
                      )
                  )
                  self.ONE_CPU = model.param_init_net.ConstantFill(
                      [], "ONE_CPU", shape=[1], value=1.0,
                  )
  
          data_parallel_model.Parallelize_GPU(
              model,
              input_builder_fun=input_builder_fun,
              forward_pass_builder_fun=model_build_fun,
              param_update_builder_fun=param_update_fun,
              devices=gpu_devices,
          )
  
          # Update the vecs
          with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, 0)):
              for num, vec in enumerate(self.vecs[:-1]):
                  model.CopyGPUToCPU("gpu_0/gpuvec_{}".format(num), vec)
  
          # Each run has same input, independent of number of gpus
          for i in range(0, 10):
              np.random.seed(2603)
              full_indices = np.random.permutation(V)[:batch_size].reshape(
                  batch_size
              )
              full_labels = full_indices[:] % batch_per_device
  
              for (j, g) in enumerate(gpu_devices):
                  st = j * batch_per_device
                  en = st + batch_per_device
                  indices = full_indices[st:en].astype(np.int32)
                  labels = full_labels[st:en].astype(np.int32)
  
                  with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, g)):
                      workspace.FeedBlob("gpu_{}/indices".format(g), indices)
                      workspace.FeedBlob("gpu_{}/label".format(g), labels)
  
              if i == 0:
                  workspace.RunNetOnce(model.param_init_net)
                  # Force vecs to be same on all runs
                  orig_vecs = [
                      np.random.rand(V, 16).astype(np.float32),
                      np.random.rand(V).astype(np.float32),
                      np.random.rand(V, 16).astype(np.float32),
                  ]
                  for vec, orig_vec in zip(self.vecs, orig_vecs):
                      workspace.FeedBlob(
                          vec,
                          orig_vec
                      )
                  for g in gpu_devices:
                      for num, orig_vec in enumerate(orig_vecs):
                          workspace.FeedBlob(
                              "gpu_{}/gpuvec_{}".format(g, num),
                              orig_vec,
                              device_option=core.DeviceOption(
                                  workspace.GpuDeviceType, g),
                          )
                  workspace.CreateNet(model.net)
  
              workspace.RunNet(model.net.Proto().name)
  
              idx = workspace.FetchBlob('gpu_0/indices')
              grad_slices = [
                  workspace.FetchBlob(
                      'gpu_{}/gpu_vec_gathered_{}_grad'.format(g, num))
                  for g in gpu_devices for num in range(2)
              ]
              for grad_slice in grad_slices:
                  # print (len(idx), len(grad_slice))
                  assert len(idx) == len(grad_slice), (
                      'Number of indices {} is not same as number of gradient '
                      'slices {}. This might lead to illegal memory access'.format(
                          len(idx), len(grad_slice)
                      )
                  )
  
      def test_sparse_shared_indices_gpu(self):
          '''
              Test that the model has same number of indices and gradient rows
              given total batchsize, independent of number of GPUs.
          '''
          V = 10000
          self.run_model(V, [0, 1])
          self.run_model(V, [0])
  
          if workspace.NumGpuDevices() >= 4:
              self.run_model(V, list(range(4)))
  
          if workspace.NumGpuDevices() >= 8:
              self.run_model(V, list(range(8)))
  
  
  if __name__ == "__main__":
      import unittest
      unittest.main()
   import unittest
  import io
  import tempfile
  import torch
  import torch.utils.show_pickle
  
  from torch.testing._internal.common_utils import IS_WINDOWS
  
  class TestShowPickle(unittest.TestCase):
  
      @unittest.skipIf(IS_WINDOWS, "Can't re-open temp file on Windows")
      def test_scripted_model(self):
          class MyCoolModule(torch.nn.Module):
              def __init__(self, weight):
                  super().__init__()
                  self.weight = weight
  
              def forward(self, x):
                  return x * self.weight
  
          m = torch.jit.script(MyCoolModule(torch.tensor([2.0])))
  
          with tempfile.NamedTemporaryFile() as tmp:
              torch.jit.save(m, tmp)
              tmp.flush()
              buf = io.StringIO()
              torch.utils.show_pickle.main(["", tmp.name + "@*/data.pkl"], output_stream=buf)
              output = buf.getvalue()
              self.assertRegex(output, "MyCoolModule")
              self.assertRegex(output, "weight")
  
  
  if __name__ == '__main__':
      unittest.main()
   from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from __future__ import unicode_literals
  
  from caffe2.python import core
  from hypothesis import given
  import caffe2.python.hypothesis_test_util as hu
  import caffe2.python.serialized_test.serialized_test_util as serial
  import hypothesis.strategies as st
  import numpy as np
  
  
  # Reference implementation from detectron/lib/utils/boxes.py
  def bbox_transform(boxes, deltas, weights=(1.0, 1.0, 1.0, 1.0)):
      """Forward transform that maps proposal boxes to predicted ground-truth
      boxes using bounding-box regression deltas. See bbox_transform_inv for a
      description of the weights argument.
      """
      if boxes.shape[0] == 0:
          return np.zeros((0, deltas.shape[1]), dtype=deltas.dtype)
  
      boxes = boxes.astype(deltas.dtype, copy=False)
  
      widths = boxes[:, 2] - boxes[:, 0] + 1.0
      heights = boxes[:, 3] - boxes[:, 1] + 1.0
      ctr_x = boxes[:, 0] + 0.5 * widths
      ctr_y = boxes[:, 1] + 0.5 * heights
  
      wx, wy, ww, wh = weights
      dx = deltas[:, 0::4] / wx
      dy = deltas[:, 1::4] / wy
      dw = deltas[:, 2::4] / ww
      dh = deltas[:, 3::4] / wh
  
      # Prevent sending too large values into np.exp()
      BBOX_XFORM_CLIP = np.log(1000. / 16.)
      dw = np.minimum(dw, BBOX_XFORM_CLIP)
      dh = np.minimum(dh, BBOX_XFORM_CLIP)
  
      pred_ctr_x = dx * widths[:, np.newaxis] + ctr_x[:, np.newaxis]
      pred_ctr_y = dy * heights[:, np.newaxis] + ctr_y[:, np.newaxis]
      pred_w = np.exp(dw) * widths[:, np.newaxis]
      pred_h = np.exp(dh) * heights[:, np.newaxis]
  
      pred_boxes = np.zeros(deltas.shape, dtype=deltas.dtype)
      # x1
      pred_boxes[:, 0::4] = pred_ctr_x - 0.5 * pred_w
      # y1
      pred_boxes[:, 1::4] = pred_ctr_y - 0.5 * pred_h
      # x2 (note: "- 1" is correct; don't be fooled by the asymmetry)
      pred_boxes[:, 2::4] = pred_ctr_x + 0.5 * pred_w - 1
      # y2 (note: "- 1" is correct; don't be fooled by the asymmetry)
      pred_boxes[:, 3::4] = pred_ctr_y + 0.5 * pred_h - 1
  
      return pred_boxes
  
  
  # Reference implementation from detectron/lib/utils/boxes.py
  def clip_tiled_boxes(boxes, im_shape):
      """Clip boxes to image boundaries. im_shape is [height, width] and boxes
      has shape (N, 4 * num_tiled_boxes)."""
      assert (
          boxes.shape[1] % 4 == 0
      ), "boxes.shape[1] is {:d}, but must be divisible by 4.".format(
          boxes.shape[1]
      )
      # x1 >= 0
      boxes[:, 0::4] = np.maximum(np.minimum(boxes[:, 0::4], im_shape[1] - 1), 0)
      # y1 >= 0
      boxes[:, 1::4] = np.maximum(np.minimum(boxes[:, 1::4], im_shape[0] - 1), 0)
      # x2 < im_shape[1]
      boxes[:, 2::4] = np.maximum(np.minimum(boxes[:, 2::4], im_shape[1] - 1), 0)
      # y2 < im_shape[0]
      boxes[:, 3::4] = np.maximum(np.minimum(boxes[:, 3::4], im_shape[0] - 1), 0)
      return boxes
  
  
  def generate_rois(roi_counts, im_dims):
      assert len(roi_counts) == len(im_dims)
      all_rois = []
      for i, num_rois in enumerate(roi_counts):
          if num_rois == 0:
              continue
          # [batch_idx, x1, y1, x2, y2]
          rois = np.random.uniform(0, im_dims[i], size=(roi_counts[i], 5)).astype(
              np.float32
          )
          rois[:, 0] = i  # batch_idx
          # Swap (x1, x2) if x1 > x2
          rois[:, 1], rois[:, 3] = (
              np.minimum(rois[:, 1], rois[:, 3]),
              np.maximum(rois[:, 1], rois[:, 3]),
          )
          # Swap (y1, y2) if y1 > y2
          rois[:, 2], rois[:, 4] = (
              np.minimum(rois[:, 2], rois[:, 4]),
              np.maximum(rois[:, 2], rois[:, 4]),
          )
          all_rois.append(rois)
      if len(all_rois) > 0:
          return np.vstack(all_rois)
      return np.empty((0, 5)).astype(np.float32)
  
  
  def bbox_transform_rotated(
      boxes,
      deltas,
      weights=(1.0, 1.0, 1.0, 1.0),
      angle_bound_on=True,
      angle_bound_lo=-90,
      angle_bound_hi=90,
  ):
      """
      Similar to bbox_transform but for rotated boxes with angle info.
      """
      if boxes.shape[0] == 0:
          return np.zeros((0, deltas.shape[1]), dtype=deltas.dtype)
  
      boxes = boxes.astype(deltas.dtype, copy=False)
  
      ctr_x = boxes[:, 0]
      ctr_y = boxes[:, 1]
      widths = boxes[:, 2]
      heights = boxes[:, 3]
      angles = boxes[:, 4]
  
      wx, wy, ww, wh = weights
      dx = deltas[:, 0::5] / wx
      dy = deltas[:, 1::5] / wy
      dw = deltas[:, 2::5] / ww
      dh = deltas[:, 3::5] / wh
      da = deltas[:, 4::5] * 180.0 / np.pi
  
      # Prevent sending too large values into np.exp()
      BBOX_XFORM_CLIP = np.log(1000. / 16.)
      dw = np.minimum(dw, BBOX_XFORM_CLIP)
      dh = np.minimum(dh, BBOX_XFORM_CLIP)
  
      pred_boxes = np.zeros(deltas.shape, dtype=deltas.dtype)
      pred_boxes[:, 0::5] = dx * widths[:, np.newaxis] + ctr_x[:, np.newaxis]
      pred_boxes[:, 1::5] = dy * heights[:, np.newaxis] + ctr_y[:, np.newaxis]
      pred_boxes[:, 2::5] = np.exp(dw) * widths[:, np.newaxis]
      pred_boxes[:, 3::5] = np.exp(dh) * heights[:, np.newaxis]
  
      pred_angle = da + angles[:, np.newaxis]
      if angle_bound_on:
          period = angle_bound_hi - angle_bound_lo
          assert period % 180 == 0
          pred_angle[np.where(pred_angle < angle_bound_lo)] += period
          pred_angle[np.where(pred_angle > angle_bound_hi)] -= period
      pred_boxes[:, 4::5] = pred_angle
  
      return pred_boxes
  
  
  def clip_tiled_boxes_rotated(boxes, im_shape, angle_thresh=1.0):
      """
      Similar to clip_tiled_boxes but for rotated boxes with angle info.
      Only clips almost horizontal boxes within angle_thresh. The rest are
      left unchanged.
      """
      assert (
          boxes.shape[1] % 5 == 0
      ), "boxes.shape[1] is {:d}, but must be divisible by 5.".format(
          boxes.shape[1]
      )
  
      (H, W) = im_shape[:2]
  
      # Filter boxes that are almost upright within angle_thresh tolerance
      idx = np.where(np.abs(boxes[:, 4::5]) <= angle_thresh)
      idx5 = idx[1] * 5
      # convert to (x1, y1, x2, y2)
      x1 = boxes[idx[0], idx5] - (boxes[idx[0], idx5 + 2] - 1) / 2.0
      y1 = boxes[idx[0], idx5 + 1] - (boxes[idx[0], idx5 + 3] - 1) / 2.0
      x2 = boxes[idx[0], idx5] + (boxes[idx[0], idx5 + 2] - 1) / 2.0
      y2 = boxes[idx[0], idx5 + 1] + (boxes[idx[0], idx5 + 3] - 1) / 2.0
      # clip
      x1 = np.maximum(np.minimum(x1, W - 1), 0)
      y1 = np.maximum(np.minimum(y1, H - 1), 0)
      x2 = np.maximum(np.minimum(x2, W - 1), 0)
      y2 = np.maximum(np.minimum(y2, H - 1), 0)
      # convert back to (xc, yc, w, h)
      boxes[idx[0], idx5] = (x1 + x2) / 2.0
      boxes[idx[0], idx5 + 1] = (y1 + y2) / 2.0
      boxes[idx[0], idx5 + 2] = x2 - x1 + 1
      boxes[idx[0], idx5 + 3] = y2 - y1 + 1
  
      return boxes
  
  
  def generate_rois_rotated(roi_counts, im_dims):
      rois = generate_rois(roi_counts, im_dims)
      # [batch_id, ctr_x, ctr_y, w, h, angle]
      rotated_rois = np.empty((rois.shape[0], 6)).astype(np.float32)
      rotated_rois[:, 0] = rois[:, 0]  # batch_id
      rotated_rois[:, 1] = (rois[:, 1] + rois[:, 3]) / 2.  # ctr_x = (x1 + x2) / 2
      rotated_rois[:, 2] = (rois[:, 2] + rois[:, 4]) / 2.  # ctr_y = (y1 + y2) / 2
      rotated_rois[:, 3] = rois[:, 3] - rois[:, 1] + 1.0  # w = x2 - x1 + 1
      rotated_rois[:, 4] = rois[:, 4] - rois[:, 2] + 1.0  # h = y2 - y1 + 1
      rotated_rois[:, 5] = np.random.uniform(-90.0, 90.0)  # angle in degrees
      return rotated_rois
  
  
  class TestBBoxTransformOp(serial.SerializedTestCase):
      @serial.given(
          num_rois=st.integers(1, 10),
          num_classes=st.integers(1, 10),
          im_dim=st.integers(100, 600),
          skip_batch_id=st.booleans(),
          rotated=st.booleans(),
          angle_bound_on=st.booleans(),
          clip_angle_thresh=st.sampled_from([-1.0, 1.0]),
          **hu.gcs_cpu_only
      )
      def test_bbox_transform(
          self,
          num_rois,
          num_classes,
          im_dim,
          skip_batch_id,
          rotated,
          angle_bound_on,
          clip_angle_thresh,
          gc,
          dc,
      ):
          """
          Test with all rois belonging to a single image per run.
          """
          rois = (
              generate_rois_rotated([num_rois], [im_dim])
              if rotated
              else generate_rois([num_rois], [im_dim])
          )
          box_dim = 5 if rotated else 4
          if skip_batch_id:
              rois = rois[:, 1:]
          deltas = np.random.randn(num_rois, box_dim * num_classes).astype(np.float32)
          im_info = np.array([im_dim, im_dim, 1.0]).astype(np.float32).reshape(1, 3)
  
          def bbox_transform_ref(rois, deltas, im_info):
              boxes = rois if rois.shape[1] == box_dim else rois[:, 1:]
              im_shape = im_info[0, 0:2]
              if rotated:
                  box_out = bbox_transform_rotated(
                      boxes, deltas, angle_bound_on=angle_bound_on
                  )
                  box_out = clip_tiled_boxes_rotated(
                      box_out, im_shape, angle_thresh=clip_angle_thresh
                  )
              else:
                  box_out = bbox_transform(boxes, deltas)
                  box_out = clip_tiled_boxes(box_out, im_shape)
              return [box_out]
  
          op = core.CreateOperator(
              "BBoxTransform",
              ["rois", "deltas", "im_info"],
              ["box_out"],
              apply_scale=False,
              correct_transform_coords=True,
              rotated=rotated,
              angle_bound_on=angle_bound_on,
              clip_angle_thresh=clip_angle_thresh,
          )
  
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=[rois, deltas, im_info],
              reference=bbox_transform_ref,
          )
  
      @given(
          roi_counts=st.lists(st.integers(0, 5), min_size=1, max_size=10),
          num_classes=st.integers(1, 10),
          rotated=st.booleans(),
          angle_bound_on=st.booleans(),
          clip_angle_thresh=st.sampled_from([-1.0, 1.0]),
          **hu.gcs_cpu_only
      )
      def test_bbox_transform_batch(
          self,
          roi_counts,
          num_classes,
          rotated,
          angle_bound_on,
          clip_angle_thresh,
          gc,
          dc,
      ):
          """
          Test with rois for multiple images in a batch
          """
          batch_size = len(roi_counts)
          total_rois = sum(roi_counts)
          im_dims = np.random.randint(100, 600, batch_size)
          rois = (
              generate_rois_rotated(roi_counts, im_dims)
              if rotated
              else generate_rois(roi_counts, im_dims)
          )
          box_dim = 5 if rotated else 4
          deltas = np.random.randn(total_rois, box_dim * num_classes).astype(np.float32)
          im_info = np.zeros((batch_size, 3)).astype(np.float32)
          im_info[:, 0] = im_dims
          im_info[:, 1] = im_dims
          im_info[:, 2] = 1.0
  
          def bbox_transform_ref(rois, deltas, im_info):
              box_out = []
              offset = 0
              for i, num_rois in enumerate(roi_counts):
                  if num_rois == 0:
                      continue
                  cur_boxes = rois[offset : offset + num_rois, 1:]
                  cur_deltas = deltas[offset : offset + num_rois]
                  im_shape = im_info[i, 0:2]
                  if rotated:
                      cur_box_out = bbox_transform_rotated(
                          cur_boxes, cur_deltas, angle_bound_on=angle_bound_on
                      )
                      cur_box_out = clip_tiled_boxes_rotated(
                          cur_box_out, im_shape, angle_thresh=clip_angle_thresh
                      )
                  else:
                      cur_box_out = bbox_transform(cur_boxes, cur_deltas)
                      cur_box_out = clip_tiled_boxes(cur_box_out, im_shape)
                  box_out.append(cur_box_out)
                  offset += num_rois
  
              if len(box_out) > 0:
                  box_out = np.vstack(box_out)
              else:
                  box_out = np.empty(deltas.shape).astype(np.float32)
              return [box_out, roi_counts]
  
          op = core.CreateOperator(
              "BBoxTransform",
              ["rois", "deltas", "im_info"],
              ["box_out", "roi_batch_splits"],
              apply_scale=False,
              correct_transform_coords=True,
              rotated=rotated,
              angle_bound_on=angle_bound_on,
              clip_angle_thresh=clip_angle_thresh,
          )
  
          self.assertReferenceChecks(
              device_option=gc,
              op=op,
              inputs=[rois, deltas, im_info],
              reference=bbox_transform_ref,
          )
  
apis: torch.testing._internal.common_cuda.initialize_cuda_context_rng torch.testing._compare_tensors_internal numpy.random.seed numpy.array torch.autograd.gradcheck torch.autograd.gradcheck.gradgradcheck
apis: torch.testing._internal.common_utils.run_tests
apis: numpy.sqrt numpy.sum numpy.allclose numpy.linalg.norm numpy.zeros numpy.testing.assert_equal numpy.var numpy.random.permutation numpy.reshape numpy.random.seed numpy.isclose numpy.random.rand numpy.round numpy.concatenate numpy.array numpy.mean
apis: torch.tensor torch.jit.save torch.utils.show_pickle.main
apis: numpy.random.uniform numpy.vstack numpy.empty numpy.zeros numpy.abs numpy.random.randn numpy.exp numpy.log numpy.maximum numpy.where numpy.random.randint numpy.array numpy.minimum