jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_horovod_allreduce_gpu_process_sets(self):
        """ Test on GPU that allreduce correctly sums if restricted to non-global process sets"""
        # Only do this test if there are GPUs available.
        if not tf.test.is_gpu_available(cuda_only=True):
            self.skipTest("No GPUs available")

        if int(os.environ.get('HOROVOD_MIXED_INSTALL', 0)):
            # Skip if compiled with CUDA but without HOROVOD_GPU_OPERATIONS.
            self.skipTest("Not compiled with HOROVOD_GPU_OPERATIONS")

        local_rank = hvd.local_rank()
        rank = hvd.rank()

        dtypes = [tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % local_rank):
                even_rank_tensor = self.random_uniform([17] * dim, -100, 100)
                even_rank_tensor = tf.cast(even_rank_tensor, dtype=dtype)
                odd_rank_tensor = self.random_uniform([17] * dim, -100, 100)
                odd_rank_tensor = tf.cast(odd_rank_tensor, dtype=dtype)
                if rank in self.even_ranks:
                    summed = hvd.allreduce(even_rank_tensor, average=False, process_set=self.even_set)
                    multiplied = even_rank_tensor * len(self.even_ranks)
                if rank in self.odd_ranks:
                    summed = hvd.allreduce(odd_rank_tensor, average=False, process_set=self.odd_set)
                    multiplied = odd_rank_tensor * len(self.odd_ranks)
                difference = summed - multiplied
                difference = tf.cast(difference, tf.int32) if dtype == tf.uint8 else difference
                max_difference = tf.reduce_max(tf.abs(difference))

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            max_process_set_size = max(len(self.even_ranks), len(self.odd_ranks))
            if max_process_set_size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif max_process_set_size < 10:
                threshold = 1e-4
            elif max_process_set_size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_cpu_process_sets(self):
        """Test on CPU that the grouped allreduce correctly sums if restricted to non-global process sets"""
        rank = hvd.rank()

        if hvd.ccl_built():
            self.skipTest("Multiple process sets currently do not support CCL.")

        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                even_rank_tensors = [tf.cast(self.random_uniform(
                    [17] * dim, -100, 100), dtype=dtype) for _ in range(5)]
                odd_rank_tensors = [tf.cast(self.random_uniform(
                    [17] * dim, -100, 100), dtype=dtype) for _ in range(5)]
                if rank in self.even_ranks:
                    summed = hvd.grouped_allreduce(even_rank_tensors, average=False, process_set=self.even_set)
                    multiplied = [t * len(self.even_ranks) for t in even_rank_tensors]
                elif rank in self.odd_ranks:
                    summed = hvd.grouped_allreduce(odd_rank_tensors, average=False, process_set=self.odd_set)
                    multiplied = [t * len(self.odd_ranks) for t in odd_rank_tensors]
            differences = [t1 - t2 for t1, t2 in zip(summed, multiplied)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            max_process_set_size = max(len(self.even_ranks), len(self.odd_ranks))
            if max_process_set_size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif max_process_set_size < 10:
                threshold = 1e-4
            elif max_process_set_size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_cpu_process_sets(self):
        """Test on CPU that the grouped allreduce correctly sums if restricted to non-global process sets"""
        rank = hvd.rank()

        if hvd.ccl_built():
            self.skipTest("Multiple process sets currently do not support CCL.")

        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                even_rank_tensors = [tf.cast(self.random_uniform(
                    [17] * dim, -100, 100), dtype=dtype) for _ in range(5)]
                odd_rank_tensors = [tf.cast(self.random_uniform(
                    [17] * dim, -100, 100), dtype=dtype) for _ in range(5)]
                if rank in self.even_ranks:
                    summed = hvd.grouped_allreduce(even_rank_tensors, average=False, process_set=self.even_set)
                    multiplied = [tensor * len(self.even_ranks) for tensor in even_rank_tensors]
                elif rank in self.odd_ranks:
                    summed = hvd.grouped_allreduce(odd_rank_tensors, average=False, process_set=self.odd_set)
                    multiplied = [tensor * len(self.odd_ranks) for tensor in odd_rank_tensors]
            differences = [t1 - t2 for t1, t2 in zip(summed, multiplied)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            max_process_set_size = max(len(self.even_ranks), len(self.odd_ranks))
            if max_process_set_size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif max_process_set_size < 10:
                threshold = 1e-4
            elif max_process_set_size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_allgather_cpu_process_sets(self):
        """Test that the allgather correctly gathers 1D, 2D, 3D tensors if restricted to non-global process sets."""
        rank = hvd.rank()

        if hvd.ccl_built():
            self.skipTest("Multiple process sets currently do not support CCL.")

        if rank in self.even_ranks:
            set_size = len(self.even_ranks)
            set_ranks = self.even_ranks
            this_set = self.even_set
        elif rank in self.odd_ranks:
            set_size = len(self.odd_ranks)
            set_ranks = self.odd_ranks
            this_set = self.odd_set

        dtypes = [tf.uint8, tf.int8, tf.uint16, tf.int16,
                  tf.int32, tf.int64, tf.float16, tf.float32,
                  tf.float64, tf.bool]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            tensor = tf.ones([17] * dim) * rank
            if dtype == tf.bool:
                tensor = tensor % 2
            tensor = tf.cast(tensor, dtype=dtype)
            with tf.device("/cpu:0"):
                gathered = hvd.allgather(tensor, process_set=this_set)

            gathered_tensor = self.evaluate(gathered)
            self.assertEqual(list(gathered_tensor.shape),
                             [17 * set_size] + [17] * (dim - 1))

            for i in range(set_size):
                rank_tensor = tf.slice(gathered_tensor,
                                       [i * 17] + [0] * (dim - 1),
                                       [17] + [-1] * (dim - 1))
                self.assertEqual(list(rank_tensor.shape), [17] * dim)
                # tf.equal() does not support tf.uint16 as of TensorFlow 1.2,
                # so need to cast rank_tensor to tf.int32.
                if dtype != tf.bool:
                    value = set_ranks[i]
                else:
                    value = set_ranks[i] % 2
                self.assertTrue(
                    self.evaluate(tf.reduce_all(
                        tf.equal(tf.cast(rank_tensor, tf.int32), value))),
                    "hvd.allgather produces incorrect gathered tensor")
//  [python]
def test_horovod_broadcast_cpu_process_sets(self):
        """Test that the broadcast correctly broadcasts 1D, 2D, 3D tensors on CPU
         if restricted to non-global process sets"""
        rank = hvd.rank()
        size = hvd.size()

        if hvd.ccl_built():
            self.skipTest("Multiple process sets currently do not support CCL.")

        # This test does not apply if there is only one worker.
        if size == 1:
            self.skipTest("Only one worker available")

        if rank in self.even_ranks:
            set_ranks = self.even_ranks
            this_set = self.even_set
        elif rank in self.odd_ranks:
            set_ranks = self.odd_ranks
            this_set = self.odd_set

        dtypes = [tf.uint8, tf.int8, tf.uint16, tf.int16,
                  tf.int32, tf.int64, tf.float16, tf.float32,
                  tf.float64, tf.bool]
        dims = [1, 2, 3]
        root_ranks = list(set_ranks)
        for dtype, dim, root_rank in itertools.product(dtypes, dims, root_ranks):
            with tf.device("/cpu:0"):
                tensor = tf.ones([17] * dim) * rank
                root_tensor = tf.ones([17] * dim) * root_rank
                if dtype == tf.bool:
                    tensor = tensor % 2
                    root_tensor = root_tensor % 2
                tensor = tf.cast(tensor, dtype=dtype)
                root_tensor = tf.cast(root_tensor, dtype=dtype)
                broadcasted_tensor = hvd.broadcast(tensor, root_rank, process_set=this_set)
            self.assertTrue(
                self.evaluate(tf.reduce_all(tf.equal(
                    tf.cast(root_tensor, tf.int32), tf.cast(broadcasted_tensor, tf.int32)))),
                "hvd.broadcast produces incorrect broadcasted tensor")
//  [python]
def test_horovod_grouped_allreduce_gpu_process_sets(self):
        """Test on GPU that the grouped allreduce correctly sums if restricted to non-global process sets"""
        if not tf.test.is_gpu_available(cuda_only=True):
            self.skipTest("No GPUs available")
        if int(os.environ.get('HOROVOD_MIXED_INSTALL', 0)):
            # Skip if compiled with CUDA but without HOROVOD_GPU_OPERATIONS.
            self.skipTest("Not compiled with HOROVOD_GPU_OPERATIONS")
        rank = hvd.rank()
        local_rank = hvd.local_rank()

        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % local_rank):
                even_rank_tensors = [tf.cast(self.random_uniform(
                    [17] * dim, -100, 100), dtype=dtype) for _ in range(5)]
                odd_rank_tensors = [tf.cast(self.random_uniform(
                    [17] * dim, -100, 100), dtype=dtype) for _ in range(5)]
                if rank in self.even_ranks:
                    summed = hvd.grouped_allreduce(even_rank_tensors, average=False, process_set=self.even_set)
                    multiplied = [t.multiply(len(self.even_ranks)) for t in even_rank_tensors]
                elif rank in self.odd_ranks:
                    summed = hvd.grouped_allreduce(odd_rank_tensors, average=False, process_set=self.odd_set)
                    multiplied = [t.multiply(len(self.odd_ranks)) for t in odd_rank_tensors]
            differences = [t1 - t2 for t1, t2 in zip(summed, multiplied)]
            differences = [t.cast(tf.int32) if dtype == tf.uint8 else diff for diff, t in zip(differences, summed)]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            max_process_set_size = max(len(self.even_ranks), len(self.odd_ranks))
            if max_process_set_size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif max_process_set_size < 10:
                threshold = 1e-4
            elif max_process_set_size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_grouped_allreduce_gpu_process_sets(self):
        """Test on GPU that the grouped allreduce correctly sums if restricted to non-global process sets"""
        if not tf.test.is_gpu_available(cuda_only=True):
            self.skipTest("No GPUs available")
        if int(os.environ.get('HOROVOD_MIXED_INSTALL', 0)):
            # Skip if compiled with CUDA but without HOROVOD_GPU_OPERATIONS.
            self.skipTest("Not compiled with HOROVOD_GPU_OPERATIONS")
        rank = hvd.rank()
        local_rank = hvd.local_rank()

        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % local_rank):
                even_rank_tensors = [tf.cast(self.random_uniform(
                    [17] * dim, -100, 100), dtype=dtype) for _ in range(5)]
                odd_rank_tensors = [tf.cast(self.random_uniform(
                    [17] * dim, -100, 100), dtype=dtype) for _ in range(5)]
                if rank in self.even_ranks:
                    summed = hvd.grouped_allreduce(even_rank_tensors, average=False, process_set=self.even_set)
                    multiplied = [tensor * len(self.even_ranks) for tensor in even_rank_tensors]
                elif rank in self.odd_ranks:
                    summed = hvd.grouped_allreduce(odd_rank_tensors, average=False, process_set=self.odd_set)
                    multiplied = [tensor * len(self.odd_ranks) for tensor in odd_rank_tensors]
            differences = [t1 - t2 for t1, t2 in zip(summed, multiplied)]
            differences = [tf.cast(diff, tf.int32) if dtype == tf.uint8 else diff for diff in differences]
            max_difference = tf.reduce_max([tf.reduce_max(tf.abs(diff)) for diff in differences])

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            max_process_set_size = max(len(self.even_ranks), len(self.odd_ranks))
            if max_process_set_size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif max_process_set_size < 10:
                threshold = 1e-4
            elif max_process_set_size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.grouped_allreduce produces incorrect results")
//  [python]
def test_horovod_allgather_gpu_process_sets(self):
        """Test that the allgather correctly gathers 1D, 2D, 3D tensors if restricted to non-global process sets."""

        if not tf.test.is_gpu_available(cuda_only=True):
            self.skipTest("No GPUs available")

        if int(os.environ.get('HOROVOD_MIXED_INSTALL', 0)):
            # Skip if compiled with CUDA but without HOROVOD_GPU_OPERATIONS.
            self.skipTest("Not compiled with HOROVOD_GPU_OPERATIONS")

        rank = hvd.rank()
        local_rank = hvd.local_rank()

        if rank in self.even_ranks:
            set_size = len(self.even_ranks)
            set_ranks = self.even_ranks
            this_set = self.even_set
        elif rank in self.odd_ranks:
            set_size = len(self.odd_ranks)
            set_ranks = self.odd_ranks
            this_set = self.odd_set

        dtypes = [tf.uint8, tf.int8, tf.uint16, tf.int16,
                  tf.int32, tf.int64, tf.float16, tf.float32,
                  tf.float64, tf.bool]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            tensor = tf.ones([17] * dim) * rank
            if dtype == tf.bool:
                tensor = tensor % 2
            tensor = tf.cast(tensor, dtype=dtype)
            with tf.device("/gpu:%d" % local_rank):
                gathered = hvd.allgather(tensor, process_set=this_set)

            gathered_tensor = self.evaluate(gathered)
            self.assertEqual(list(gathered_tensor.shape),
                             [17 * set_size] + [17] * (dim - 1))

            for i in range(set_size):
                rank_tensor = tf.slice(gathered_tensor,
                                       [i * 17] + [0] * (dim - 1),
                                       [17] + [-1] * (dim - 1))
                self.assertEqual(list(rank_tensor.shape), [17] * dim)
                # tf.equal() does not support tf.uint16 as of TensorFlow 1.2,
                # so need to cast rank_tensor to tf.int32.
                if dtype != tf.bool:
                    value = set_ranks[i]
                else:
                    value = set_ranks[i] % 2
                self.assertTrue(
                    self.evaluate(tf.reduce_all(
                        tf.equal(tf.cast(rank_tensor, tf.int32), value))),
                    "hvd.allgather produces incorrect gathered tensor")
//  [python]
def test_horovod_broadcast_grad_cpu_process_sets(self):
        """Test the correctness of the broadcast gradient on CPU if restricted to non-global process sets."""
        rank = hvd.rank()
        size = hvd.size()

        if hvd.ccl_built():
            self.skipTest("Multiple process sets currently do not support CCL.")

        # This test does not apply if there is only one worker.
        if size == 1:
            self.skipTest("Only one worker available")

        if rank in self.even_ranks:
            set_ranks = self.even_ranks
            this_set = self.even_set
        elif rank in self.odd_ranks:
            set_ranks = self.odd_ranks
            this_set = self.odd_set

        # As of TensorFlow v1.9, gradients are not supported on
        # integer tensors
        dtypes = [tf.float32, tf.float64]
        dims = [1, 2, 3]
        root_ranks = list(set_ranks)
        for dtype, dim, root_rank in itertools.product(dtypes, dims, root_ranks):
            with tf.device("/cpu:0"):
                if _executing_eagerly():
                    tensor = self.tfe.Variable(tf.ones([5] * dim) * rank)
                else:
                    tensor = tf.ones([5] * dim) * rank
                if dtype == tf.bool:
                    tensor = tensor % 2
                if _executing_eagerly():
                    with tf.GradientTape() as tape:
                        tensor = tf.cast(tensor, dtype=dtype)
                        broadcasted_tensor = hvd.broadcast(tensor, root_rank,
                                                           process_set=this_set)
                    grad_out = tape.gradient(broadcasted_tensor, tensor)
                else:
                    tensor = tf.cast(tensor, dtype=dtype)
                    broadcasted_tensor = hvd.broadcast(tensor, root_rank,
                                                       process_set=this_set)
                    grad_ys = tf.ones([5] * dim)
                    grad = tf.gradients(broadcasted_tensor, tensor, grad_ys)[0]
                    grad_out = self.evaluate(grad)

            c = 1 if rank == root_rank else 0
            expected = np.ones([5] * dim) * c
            err = np.linalg.norm(expected - grad_out)
            self.assertLess(err, 0.00000001,
                            "gradient %s differs from expected %s, "
                            "error: %s" % (grad_out, expected, str(err)))
//  [python]
def __init__(self, initial_state=0, dtype=tf.int64, scope='TFEnviroment'):
    self._dtype = dtype
    self._scope = scope
    self._initial_state = t.cast(initial_state, dtype=self._dtype)
    observation_spec = specs.TensorSpec([1], self._dtype, 'observation')
    action_spec = specs.BoundedTensorSpec([], tf.int32, minimum=0, maximum=10)
    time_step_spec = ts.time_step_spec(observation_spec)
    super(TFEnvironmentMock, self).__init__(time_step_spec, action_spec)
    self._state = common.create_variable(
        'state', initial_state, dtype=self._dtype
    )
    self.steps = common.create_variable('steps', 0)
    self.episodes = common.create_variable('episodes', 0)
    self.resets = common.create_variable('resets', 0)
//  [python]
def __init__(self, initial_state=0, dtype=tf.int64, scope='TFEnviroment'):
    self._dtype = dtype
    self._scope = scope
    self._initial_state = tf.cast(initial_state, dtype=self._dtype)
    observation_spec = specs.TensorSpec([1], self._dtype, 'observation')
    action_spec = specs.BoundedTensorSpec([], tf.int32, minimum=0, maximum=10)
    time_step_spec = ts.time_step_spec(observation_spec)
    super(TFEnvironmentMock, self).__init__(time_step_spec, action_spec)
    self._state = common.create_variable(
        'state', initial_state, dtype=self._dtype
    )
    self.steps = common.create_variable('steps', 0)
    self.episodes = common.create_variable('episodes', 0)
    self.resets = common.create_variable('resets', 0)
//  [python]
def process_metrics(self, metrics: List[Any],
                      labels: Mapping[str, tf.Tensor],
                      model_outputs: tf.Tensor):
    """Processes and updates metrics.

    Called when using custom training loop API.

    Args:
      metrics: a nested structure of metrics objects. The return of function
        self.build_metrics.
      labels: a tensor or a nested structure of tensors.
      model_outputs: a tensor or a nested structure of tensors. For example,
        output of the keras model built by self.build_model.
    """
    num_classes = self.task_config.train_data.num_classes
    class_targets = tf.reshape(labels['class_target'], [-1, num_classes])
    model_outputs = tf.reshape(model_outputs, [-1, num_classes])
    sample_weight = tf.to_float(tf.reshape(labels['instances_mask'], [-1]))
    for metric in metrics:
      metric.update_state(
          y_true=class_targets,
          y_pred=model_outputs,
          sample_weight=sample_weight)
//  [python]
def process_metrics(self, metrics: List[Any],
                      labels: Mapping[str, tf.Tensor],
                      model_outputs: tf.Tensor):
    """Processes and updates metrics.

    Called when using custom training loop API.

    Args:
      metrics: a nested structure of metrics objects. The return of function
        self.build_metrics.
      labels: a tensor or a nested structure of tensors.
      model_outputs: a tensor or a nested structure of tensors. For example,
        output of the keras model built by self.build_model.
    """
    num_classes = self.task_config.train_data.num_classes
    class_targets = tf.reshape(labels['class_target'], [-1, num_classes])
    model_outputs = tf.reshape(model_outputs, [-1, num_classes])
    sample_weight = tf.cast(tf.reshape(labels['instances_mask'], [-1]),
                            tf.float32)
    for metric in metrics:
      metric.update_state(
          y_true=class_targets,
          y_pred=model_outputs,
          sample_weight=sample_weight)
//  [python]
def test_horovod_reducescatter_cpu_process_sets(self):
        """Test on CPU that the reducescatter correctly sums or averages and scatters 1D, 2D, 3D tensors
        if restricted to non-global process sets."""
        if hvd.ccl_built():
            self.skipTest("Reducescatter is not supported yet with oneCCL operations.")
        if _is_mac and hvd.gloo_built() and not hvd.mpi_built():
            self.skipTest("ReducescatterGloo is not supported on macOS")
        rank = hvd.rank()

        if rank in self.even_ranks:
            this_set = self.even_set
        else:
            this_set = self.odd_set
        process_set_size = this_set.size()
        process_set_rank = this_set.rank()

        dtypes = self.filter_supported_types([tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for red_op, dtype, dim in itertools.product([hvd.Sum, hvd.Average], dtypes, dims):
            with tf.device("/cpu:0"):
                even_rank_tensor = self.random_uniform([process_set_size * 4] * dim, -100, 100, dtype=dtype)
                odd_rank_tensor = self.random_uniform([process_set_size * 4] * dim, -100, 100, dtype=dtype)
                if rank in self.even_ranks:
                    tensor = even_rank_tensor
                else:
                    tensor = odd_rank_tensor
                reduced = hvd.reducescatter(tensor, op=red_op, process_set=this_set)
            if red_op == hvd.Sum:
                expected = tf.cast(tensor[process_set_rank * 4:(process_set_rank + 1) * 4] * process_set_size,
                                   reduced.dtype)
            elif red_op == hvd.Average:
                expected = tf.cast(tensor[process_set_rank * 4:(process_set_rank + 1) * 4], reduced.dtype)
            max_difference = tf.reduce_max(tf.abs(reduced - expected))

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            if dtype == tf.float16:
                threshold = .5
            elif dtype in [tf.int32, tf.int64]:
                threshold = 0
            elif process_set_size < 10:
                threshold = 1e-4
            elif process_set_size < 15:
                threshold = 5e-4
            else:
                break

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold,
                            "hvd.reducescatter produces incorrect results")
//  [python]
def build_losses(self,
                   labels: Mapping[str, tf.Tensor],
                   model_outputs: tf.Tensor,
                   aux_losses: Optional[Any] = None):
    """Sparse categorical cross entropy loss.

    Args:
      labels: labels dictionary contains multi-hot "class_target" and
        "sample_weights".
      model_outputs: Output logits of the classifier.
      aux_losses: auxiliarly loss tensors, i.e. `losses` in keras.Model.

    Returns:
      The dictionary contains total loss tensors.
    """
    losses_config = self.task_config.losses

    # in shape [B, N]
    xent_loss_fn = tf_keras.losses.BinaryCrossentropy(
        reduction=tf_keras.losses.Reduction.NONE,
        from_logits=False,
        label_smoothing=losses_config.label_smoothing)
    class_targets = labels['class_target']
    sample_weight = tf.cast(labels['instances_mask'], tf.float32)
    model_loss = xent_loss_fn(y_true=class_targets,
                              y_pred=model_outputs,
                              sample_weight=sample_weight)
    if self._is_multilabel():
      # Re-scale the binary cross entropy loss with num_classes to make it
      # comparable with categorical cross entropy in scale. This enables to use
      # the same learning rate for different losses.
      model_loss *= self._get_num_classes()

    model_loss = tf.reduce_sum(model_loss) / (
        tf.reduce_sum(sample_weight) + 1e-6)

    if aux_losses:
      regularization_loss = tf.add_n(aux_losses)
    else:
      regularization_loss = 0.0

    total_loss = model_loss + regularization_loss
    logs = {
        'model_loss': model_loss,
        'regularization_loss': regularization_loss,
        # Include model predictions and corresponding boxes for the eval.
        'predictions': model_outputs,
        # Register total loss.
        self.loss: total_loss,
    }
    return logs
//  [python]
def test_horovod_grouped_allgather_cpu_process_sets(self):
        """Test that the grouped allgather correctly gathers 1D, 2D, 3D tensors
        if restricted to non-global process sets."""
        rank = hvd.rank()

        if hvd.ccl_built():
            self.skipTest("Multiple process sets currently do not support CCL.")

        if rank in self.even_ranks:
            set_size = len(self.even_ranks)
            set_ranks = self.even_ranks
            this_set = self.even_set
        elif rank in self.odd_ranks:
            set_size = len(self.odd_ranks)
            set_ranks = self.odd_ranks
            this_set = self.odd_set

        dtypes = [tf.uint8, tf.int8, tf.uint16, tf.int16,
                  tf.int32, tf.int64, tf.float16, tf.float32,
                  tf.float64, tf.bool]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            tensors = [tf.ones([17] * dim) * rank for _ in range(5)]
            if dtype == tf.bool:
                tensors = [tensor % 2 for tensor in tensors]
            tensors = [tf.cast(tensor, dtype=dtype) for tensor in tensors]
            with tf.device("/cpu:0"):
                gathered = hvd.grouped_allgather(tensors, process_set=this_set)

            gathered_tensors = self.evaluate(gathered)
            for gathered_tensor in gathered_tensors:
                self.assertEqual(list(gathered_tensor.shape),
                                 [17 * set_size] + [17] * (dim - 1))

            for i in range(set_size):
                rank_tensors = [tf.slice(gathered_tensor,
                                         [i * 17] + [0] * (dim - 1),
                                         [17] + [-1] * (dim - 1))
                                for gathered_tensor in gathered_tensors]
                self.assertEqual([rank_tensor.shape for rank_tensor in rank_tensors], len(tensors) * [[17] * dim])
                # tf.equal() does not support tf.uint16 as of TensorFlow 1.2,
                # so need to cast rank_tensor to tf.int32.
                if dtype != tf.bool:
                    value = set_ranks[i]
                else:
                    value = set_ranks[i] % 2
                self.assertTrue(all(self.evaluate(tf.reduce_all(
                    tf.equal(tf.cast(rank_tensor, tf.int32), value))) for rank_tensor in rank_tensors),
                    "hvd.grouped_allgather produces incorrect gathered tensor")
//  [python]
def roi_pooling(self, feature_maps, rois, img_shape, scope):
        '''
        Here use roi warping as roi_pooling

        :param featuremaps_dict: feature map to crop
        :param rois: shape is [-1, 4]. [x1, y1, x2, y2]
        :return:
        '''

        with tf.variable_scope('ROI_Warping_'+scope):
            img_h, img_w = tf.to_float(img_shape[1]), tf.to_float(img_shape[2])
            N = tf.shape(rois)[0]
            x1, y1, x2, y2 = tf.unstack(rois, axis=1)

            normalized_x1 = x1 / img_w
            normalized_x2 = x2 / img_w
            normalized_y1 = y1 / img_h
            normalized_y2 = y2 / img_h

            normalized_rois = tf.transpose(
                tf.stack([normalized_y1, normalized_x1, normalized_y2, normalized_x2]), name='get_normalized_rois')

            normalized_rois = tf.stop_gradient(normalized_rois)

            cropped_roi_features = tf.image.crop_and_resize(feature_maps, normalized_rois,
                                                            box_ind=tf.zeros(shape=[N, ],
                                                                             dtype=tf.int32),
                                                            crop_size=[cfgs.ROI_SIZE, cfgs.ROI_SIZE],
                                                            name='CROP_AND_RESIZE'
                                                            )
            roi_features = slim.max_pool2d(cropped_roi_features,
                                           [cfgs.ROI_POOL_KERNEL_SIZE, cfgs.ROI_POOL_KERNEL_SIZE],
                                           stride=cfgs.ROI_POOL_KERNEL_SIZE)

        return roi_features
//  [python]
def roi_pooling(self, feature_maps, rois, img_shape, scope):
        '''
        Here use roi warping as roi_pooling

        :param featuremaps_dict: feature map to crop
        :param rois: shape is [-1, 4]. [x1, y1, x2, y2]
        :return:
        '''

        with tf.variable_scope('ROI_Warping_'+scope):
            img_h, img_w = tf.cast(img_shape[1], tf.float32), tf.cast(img_shape[2], tf.float32)
            N = tf.shape(rois)[0]
            x1, y1, x2, y2 = tf.unstack(rois, axis=1)

            normalized_x1 = x1 / img_w
            normalized_x2 = x2 / img_w
            normalized_y1 = y1 / img_h
            normalized_y2 = y2 / img_h

            normalized_rois = tf.transpose(
                tf.stack([normalized_y1, normalized_x1, normalized_y2, normalized_x2]), name='get_normalized_rois')

            normalized_rois = tf.stop_gradient(normalized_rois)

            cropped_roi_features = tf.image.crop_and_resize(feature_maps, normalized_rois,
                                                            box_ind=tf.zeros(shape=[N, ],
                                                                             dtype=tf.int32),
                                                            crop_size=[cfgs.ROI_SIZE, cfgs.ROI_SIZE],
                                                            name='CROP_AND_RESIZE'
                                                            )
            roi_features = slim.max_pool2d(cropped_roi_features,
                                           [cfgs.ROI_POOL_KERNEL_SIZE, cfgs.ROI_POOL_KERNEL_SIZE],
                                           stride=cfgs.ROI_POOL_KERNEL_SIZE)

        return roi_features
//  [python]
def _decode_masks(self, parsed_tensors: TensorDict) -> TensorDict:
    """Decode a set of PNG masks to the tf.float32 tensors."""

    def _decode_png_mask(png_bytes):
      mask = tf.squeeze(
          tf.io.decode_png(png_bytes, channels=1, dtype=tf.uint8), axis=-1)
      mask = tf.to_float(mask)
      mask.set_shape([None, None])
      return mask

    height = parsed_tensors['image/height']
    width = parsed_tensors['image/width']
    masks = parsed_tensors['image/object/mask']
    masks = tf.cond(
        pred=tf.greater(tf.size(input=masks), 0),
        true_fn=lambda: tf.map_fn(_decode_png_mask, masks, dtype=tf.float32),
        false_fn=lambda: tf.zeros([0, height, width], dtype=tf.float32))
    return {'groundtruth_instance_masks': masks}
//  [python]
def _decode_masks(self, parsed_tensors: TensorDict) -> TensorDict:
    """Decode a set of PNG masks to the tf.float32 tensors."""

    def _decode_png_mask(png_bytes):
      mask = tf.squeeze(
          tf.io.decode_png(png_bytes, channels=1, dtype=tf.uint8), axis=-1)
      mask = tf.cast(mask, dtype=tf.float32)
      mask.set_shape([None, None])
      return mask

    height = parsed_tensors['image/height']
    width = parsed_tensors['image/width']
    masks = parsed_tensors['image/object/mask']
    masks = tf.cond(
        pred=tf.greater(tf.size(input=masks), 0),
        true_fn=lambda: tf.map_fn(_decode_png_mask, masks, dtype=tf.float32),
        false_fn=lambda: tf.zeros([0, height, width], dtype=tf.float32))
    return {'groundtruth_instance_masks': masks}
//  [python]
def call(self, encoder_outputs, doc_attention_mask):
    num_docs = tf_utils.get_shape_list(encoder_outputs, expected_rank=[4])[1]
    cls_embeddings = encoder_outputs[:, :, 0, :]
    key = self._key_dense(cls_embeddings)
    query = self._query_dense(cls_embeddings)
    doc_attention_mask = tf.cast(doc_attention_mask, tf.float32)

    key = tf.einsum("BANH,BA->BANH", key, doc_attention_mask)
    query = tf.einsum("BANH,BA->BANH", query, doc_attention_mask)
    attention_matrix = tf.einsum("BXNH,BYNH->BNXY", query, key)
    mask = tf.ones([num_docs, num_docs])
    mask = tf.linalg.set_diag(mask, tf.zeros(num_docs))
    attention_matrix = tf.einsum("BNXY,XY->BNXY", attention_matrix, mask)
    doc_attention_probs = tf.einsum("BNAY->BNA", attention_matrix)
    doc_attention_probs = tf.einsum("BNA->BA", doc_attention_probs)
    infadder = (1.0 - doc_attention_mask) * -100000.0
    return tf.nn.softmax(doc_attention_probs + infadder)
//  [python]
def _normalize_inputs(self, inputs):
        """Normalize inputs"""
        inputs = tf.expand_dims(tf.cast(inputs, tf.float32), -1)

        # Center around zero
        inputs = tf.divide(inputs, 127.5) - 1

        inputs = tf.image.resize_images(inputs, (388, 388))

        return tf.image.resize_image_with_crop_or_pad(inputs, 572, 572)
//  [python]
def call(
        self,
        input_tensor: tf.Tensor,
        attention_mask: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        batch_size = shape_list(input_tensor)[0]
        mixed_query_layer = self.query(inputs=input_tensor)
        mixed_key_layer = self.key(inputs=input_tensor)
        mixed_value_layer = self.value(inputs=input_tensor)
        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.to_float(self.sqrt_att_head_size)
        attention_scores = tf.divide(attention_scores, dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFAlbertModel call() function)
            attention_scores = tf.add(attention_scores, attention_mask)

        # Normalize the attention scores to probabilities.
        attention_probs = stable_softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.attention_dropout(inputs=attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = tf.multiply(attention_probs, head_mask)

        context_layer = tf.matmul(attention_probs, value_layer)
        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])

        # (batch_size, seq_len_q, all_head_size)
        context_layer = tf.reshape(tensor=context_layer, shape=(batch_size, -1, self.all_head_size))
        self_outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
        hidden_states = self_outputs[0]
        hidden_states = self.dense(inputs=hidden_states)
        hidden_states = self.output_dropout(inputs=hidden_states, training=training)
        attention_output = self.LayerNorm(inputs=hidden_states + input_tensor)

        # add attentions if we output them
        outputs = (attention_output,) + self_outputs[1:]

        return outputs
//  [python]
def call(
        self,
        input_tensor: tf.Tensor,
        attention_mask: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        batch_size = shape_list(input_tensor)[0]
        mixed_query_layer = self.query(inputs=input_tensor)
        mixed_key_layer = self.key(inputs=input_tensor)
        mixed_value_layer = self.value(inputs=input_tensor)
        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # (batch size, num_heads, seq_len_q, seq_len_k)
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)
        attention_scores = tf.divide(attention_scores, dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFAlbertModel call() function)
            attention_scores = tf.add(attention_scores, attention_mask)

        # Normalize the attention scores to probabilities.
        attention_probs = stable_softmax(logits=attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.attention_dropout(inputs=attention_probs, training=training)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = tf.multiply(attention_probs, head_mask)

        context_layer = tf.matmul(attention_probs, value_layer)
        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])

        # (batch_size, seq_len_q, all_head_size)
        context_layer = tf.reshape(tensor=context_layer, shape=(batch_size, -1, self.all_head_size))
        self_outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
        hidden_states = self_outputs[0]
        hidden_states = self.dense(inputs=hidden_states)
        hidden_states = self.output_dropout(inputs=hidden_states, training=training)
        attention_output = self.LayerNorm(inputs=hidden_states + input_tensor)

        # add attentions if we output them
        outputs = (attention_output,) + self_outputs[1:]

        return outputs
//  [python]
def serve(self, images: tf.Tensor) -> dict[str, tf.Tensor]:
    """Cast image to float and run inference.

    Args:
      images: uint8 Tensor of shape [batch_size, None, None, 3]

    Returns:
      Tensor holding classification output logits.
    """
    # Skip image preprocessing when input_type is tflite so it is compatible
    # with TFLite quantization.
    image_info = None
    if self._input_type != 'tflite':
      with tf.device('cpu:0'):
        images = tf.to_float(images)

        images_spec = tf.TensorSpec(
            shape=self._input_image_size + [3], dtype=tf.float32)
        image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)

        images, image_info = tf.nest.map_structure(
            tf.identity,
            tf.map_fn(
                self._build_inputs,
                elems=images,
                fn_output_signature=(images_spec, image_info_spec),
                parallel_iterations=32))

    outputs = self.inference_step(images)[-1]
    outputs = {
        'detection_boxes': outputs['detection_boxes'],
        'detection_scores': outputs['detection_scores'],
        'detection_classes': outputs['detection_classes'],
        'num_detections': outputs['num_detections']
    }
    if image_info is not None:
      outputs['detection_boxes'] = outputs['detection_boxes'] * tf.expand_dims(
          tf.concat([
              image_info[:, 1:2, 0], image_info[:, 1:2, 1],
              image_info[:, 1:2, 0], image_info[:, 1:2, 1]
          ],
                    axis=1),
          axis=1)

      outputs.update({'image_info': image_info})

    return outputs
//  [python]
def serve(self, images: tf.Tensor) -> dict[str, tf.Tensor]:
    """Cast image to float and run inference.

    Args:
      images: uint8 Tensor of shape [batch_size, None, None, 3]

    Returns:
      Tensor holding classification output logits.
    """
    # Skip image preprocessing when input_type is tflite so it is compatible
    # with TFLite quantization.
    image_info = None
    if self._input_type != 'tflite':
      with tf.device('cpu:0'):
        images = tf.cast(images, dtype=tf.float32)

        images_spec = tf.TensorSpec(
            shape=self._input_image_size + [3], dtype=tf.float32)
        image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)

        images, image_info = tf.nest.map_structure(
            tf.identity,
            tf.map_fn(
                self._build_inputs,
                elems=images,
                fn_output_signature=(images_spec, image_info_spec),
                parallel_iterations=32))

    outputs = self.inference_step(images)[-1]
    outputs = {
        'detection_boxes': outputs['detection_boxes'],
        'detection_scores': outputs['detection_scores'],
        'detection_classes': outputs['detection_classes'],
        'num_detections': outputs['num_detections']
    }
    if image_info is not None:
      outputs['detection_boxes'] = outputs['detection_boxes'] * tf.expand_dims(
          tf.concat([
              image_info[:, 1:2, 0], image_info[:, 1:2, 1],
              image_info[:, 1:2, 0], image_info[:, 1:2, 1]
          ],
                    axis=1),
          axis=1)

      outputs.update({'image_info': image_info})

    return outputs
//  [python]
def build_model(self):

        self.X = tf.placeholder(tf.float32, [None, self.dims[0]])

        self.Y = tf.placeholder(tf.float32, [None, self.max_length])

        self.N_sample = tf.placeholder(dtype=tf.float32, shape=[1])


        lstm_cell = tf.contrib.rnn.BasicLSTMCell(num_units=self.rnn_hidden_size)

        self.segment_ids = tf.placeholder(tf.float32, [self.max_length])

        self.rnn_dim_feature = np.shape(self.data_X)[-1]

        self.rnn_input = tf.constant(value=self.data_X, dtype=tf.float32, shape=[self.num_example,self.max_length, self.rnn_dim_feature]) #

        state = lstm_cell.zero_state(self.num_example, dtype=tf.float32)

        xx = tf.unstack(self.rnn_input, self.max_length, 1)

        rnn_output, state = rnn.static_rnn(lstm_cell, xx, initial_state=state)

        self.train_loss = 0
        self.YY = {}

        for step in range(self.max_length):


            self.YY[str(step)] = tf.slice(self.Y, [0, step], [-1, 1])

            ss = tf.slice(rnn_output, [step, 0, 0], [1, -1, -1])

            self.embed[str(step)] = self.deepset(tf.squeeze(ss))


            state_embedding = tf.tile(self.embed[str(step)], [tf.to_int32(self.N_sample[0]), 1])



            self.bo_net(self.X,state_embedding, step)


            self.train_loss += self.ml_primal[str(step)]




        self.opt = tf.train.AdamOptimizer(learning_rate=self.lr).minimize(self.train_loss)