def _is_finite(arg1, *args):
  """Checks if the supplied tensors are finite.

  Args:
    arg1: A numeric `Tensor`.
    *args: (Optional) Other `Tensors` to check for finiteness.

  Returns:
    is_finite: Scalar boolean `Tensor` indicating whether all the supplied
      tensors are finite.
  """
  finite = tf.reduce_all(tf.is_finite(arg1))
  for arg in args:
    finite = finite & tf.reduce_all(tf.is_finite(arg))
  return finite

def _has_enough_pixels_of_each_object_in_first_frame(
    label, decoder_output_stride):
  """Checks if for each object (incl. background) enough pixels are visible.

  During test time, we will usually not see a reference frame in which only
  very few pixels of one object are visible. These cases can be problematic
  during training, especially if more than the 1-nearest neighbor is used.
  That's why this function can be used to detect and filter these cases.

  Args:
    label: Label tensor of shape [num_frames, height, width, 1].
    decoder_output_stride: Integer, the stride of the decoder output.

  Returns:
    Boolean, whether the labels have enough pixels of each object in the first
      frame.
  """
  h, w = train_utils.resolve_shape(label)[1:3]
  h_sub = model.scale_dimension(h, 1.0 / decoder_output_stride)
  w_sub = model.scale_dimension(w, 1.0 / decoder_output_stride)
  label_downscaled = tf.squeeze(
      tf.image.resize_nearest_neighbor(label[0, tf.newaxis], [h_sub, w_sub],
                                       align_corners=True), axis=0)
  _, _, counts = tf.unique_with_counts(
      tf.reshape(label_downscaled, [-1]))
  has_enough_pixels_per_object = tf.reduce_all(
      tf.greater_equal(counts, MIN_LABEL_COUNT))
  return has_enough_pixels_per_object

    def loop_continue_criterion(self, *args) -> tf.Tensor:
        """Decide whether to break out of the while loop.

        The criterion for stopping the loop is that either all hypotheses are
        finished or a maximum number of steps has been reached. Here the number
        of steps is the number of steps of the underlying decoder minus one,
        because this function is evaluated after the decoder step has been
        called and its step has been incremented. This is caused by the fact
        that we call the decoder body function at the end of the beam body
        function. (And that, in turn, is to support ensembling.)

        Arguments:
            args: A ``BeamSearchLoopState`` instance.

        Returns:
            A scalar boolean ``Tensor``.
        """
        loop_state = BeamSearchLoopState(*args)

        beam_step = loop_state.decoder_loop_state.feedables.step - 1
        finished = loop_state.search_state.finished

        max_step_cond = tf.less(beam_step, self.max_steps)
        unfinished_cond = tf.logical_not(tf.reduce_all(finished))

        return tf.logical_and(max_step_cond, unfinished_cond)

    def loop_fn(time, cell_output, cell_state, loop_state):
        emit_output = cell_output  # == None for time == 0
        
        if cell_output is None:  # time == 0
            next_cell_state = initial_state
        else:
            next_cell_state = cell_state
        
        elements_finished = (time >= sequence_length)
        
        finished = tf.reduce_all(elements_finished)
        
        next_input = tf.cond(
            finished,
            lambda: tf.zeros([batch_size, input_size], dtype=tf.float32),
            lambda: inputs_ta.read(time))

        next_target = tf.cond(
            finished,
            lambda: tf.zeros([batch_size, target_size], dtype=tf.float32),
            lambda: targets_ta.read(time))

        
        if loop_state is None:
            next_input = tf.expand_dims(next_input, 1)
            next_input = tf.pad(next_input, [[0,0], [window-1, 0], [0,0]])
        else:
            next_input = cat_hist(loop_state, next_input, 1)
        
        next_loop_state = next_input

        return (elements_finished, RnnHistInputTuple(next_input, next_target), next_cell_state, emit_output, next_loop_state)

    def test_horovod_broadcast(self):
        """Test that the broadcast correctly broadcasts 1D, 2D, 3D tensors."""
        hvd.init()
        rank = hvd.rank()
        size = hvd.size()

        # This test does not apply if there is only one worker.
        if size == 1:
            return

        with self.test_session() as session:
            dtypes = [tf.uint8, tf.int8, tf.uint16, tf.int16,
                      tf.int32, tf.int64, tf.float32, tf.float64,
                      tf.bool]
            dims = [1, 2, 3]
            root_ranks = list(range(size))
            for dtype, dim, root_rank in itertools.product(dtypes, dims, root_ranks):
                try:
                    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)
                    self.assertTrue(
                        session.run(tf.reduce_all(tf.equal(
                            tf.cast(root_tensor, tf.int32), tf.cast(broadcasted_tensor, tf.int32)))),
                        "hvd.broadcast produces incorrect broadcasted tensor")
                except Exception:
                    import traceback
                    traceback.print_exc()

def loop_fn_transition(time, previous_output, previous_state, previous_loop_state):
    def get_next_input():
        # dot product between previous ouput and weights, then + biases
        output_logits = tf.add(tf.matmul(previous_output, W), b)
        # Logits simply means that the function operates on the unscaled output of
        # earlier layers and that the relative scale to understand the units is linear.
        # It means, in particular, the sum of the inputs may not equal 1, that the values are not probabilities
        # (you might have an input of 5).
        # prediction value at current time step

        # Returns the index with the largest value across axes of a tensor.
        # Attention focusing
        prediction = tf.argmax(output_logits, axis=1)
        # embed prediction for the next input
        next_input = tf.nn.embedding_lookup(embeddings, prediction)
        return next_input

    elements_finished = (time >= decoder_length)  # this operation produces boolean tensor of [batch_size]
    # defining if corresponding sequence has ended
    # Computes the "logical and" of elements across dimensions of a tensor.
    finished = tf.reduce_all(elements_finished)  # -> boolean scalar
    # Return either fn1() or fn2() based on the boolean predicate pred.
    input = tf.cond(finished, lambda: pad_step_embedded, get_next_input)

    # set previous to current
    state = previous_state
    output = previous_output
    loop_state = None

    return (elements_finished, input, state,
            output, loop_state)

def new_mean_squared(grad_vec, decay, ms):
  """Calculates the new accumulated mean squared of the gradient.

  Args:
    grad_vec: the vector for the current gradient
    decay: the decay term
    ms: the previous mean_squared value

  Returns:
    the new mean_squared value
  """
  decay_size = decay.get_shape().num_elements()
  decay_check_ops = [
      tf.assert_less_equal(decay, 1., summarize=decay_size),
      tf.assert_greater_equal(decay, 0., summarize=decay_size)]

  with tf.control_dependencies(decay_check_ops):
    grad_squared = tf.square(grad_vec)

  # If the previous mean_squared is the 0 vector, don't use the decay and just
  # return the full grad_squared. This should only happen on the first timestep.
  decay = tf.cond(tf.reduce_all(tf.equal(ms, 0.)),
                  lambda: tf.zeros_like(decay, dtype=tf.float32), lambda: decay)

  # Update the running average of squared gradients.
  epsilon = 1e-12
  return (1. - decay) * (grad_squared + epsilon) + decay * ms

 def testUniformSamplePdf(self):
   with self.test_session():
     a = 10.0
     b = [11.0, 100.0]
     uniform = tf.contrib.distributions.Uniform(a, b)
     self.assertTrue(tf.reduce_all(uniform.pdf(uniform.sample_n(10)) > 0).eval(
     ))

    def _verify_compatible_image_shapes(img1, img2):
        """
        Checks if two image tensors are compatible for applying SSIM or PSNR.
        This function checks if two sets of images have ranks at least 3, and if the
        last three dimensions match.
        Args:
        img1: Tensor containing the first image batch.
        img2: Tensor containing the second image batch.
        Returns:
        A tuple containing: the first tensor shape, the second tensor shape, and a
        list of control_flow_ops.Assert() ops implementing the checks.
        Raises:
        ValueError: When static shape check fails.
        """
        shape1 = img1.get_shape().with_rank_at_least(3)
        shape2 = img2.get_shape().with_rank_at_least(3)
        shape1[-3:].assert_is_compatible_with(shape2[-3:])

        if shape1.ndims is not None and shape2.ndims is not None:
            for dim1, dim2 in zip(reversed(shape1[:-3]), reversed(shape2[:-3])):
                if not (dim1 == 1 or dim2 == 1 or dim1.is_compatible_with(dim2)):
                    raise ValueError('Two images are not compatible: %s and %s' % (shape1, shape2))

        # Now assign shape tensors.
        shape1, shape2 = tf.shape_n([img1, img2])

        # TODO(sjhwang): Check if shape1[:-3] and shape2[:-3] are broadcastable.
        checks = []
        checks.append(tf.Assert(tf.greater_equal(tf.size(shape1), 3),
                                [shape1, shape2], summarize=10))
        checks.append(tf.Assert(tf.reduce_all(tf.equal(shape1[-3:], shape2[-3:])),
                                [shape1, shape2], summarize=10))

        return shape1, shape2, checks

 def testUniformSamplePdf(self):
   a = 10.0
   b = [11.0, 100.0]
   uniform = uniform_lib.Uniform(a, b)
   self.assertTrue(
       self.evaluate(
           tf.reduce_all(uniform.prob(uniform.sample(10)) > 0)))

    def loss_fn(inputs, targets, input_sequence_length, output_sequence_length):
      """Creates the loss and the exports."""
      logits = model(
          inputs, targets, input_sequence_length, output_sequence_length)
      targets = tf.cast(targets, tf.int32)
      sq_sz_out_max = targets.shape[0].value

      # Create a mask to ignore accuracy on buffer characters.
      sequence_sizes = tf.cast(output_sequence_length, tf.float32)
      lengths_transposed = tf.expand_dims(sequence_sizes, 1)
      range_row = tf.expand_dims(
          tf.range(0, sq_sz_out_max, 1, dtype=tf.float32), 0)
      mask = tf.cast(tf.transpose(tf.less(range_row, lengths_transposed)),
                     tf.float32)

      # Compute token accuracy and solved.
      correct = tf.equal(tf.argmax(logits, 2), tf.argmax(targets, 2))
      solved = tf.reduce_all(tf.boolean_mask(correct, tf.squeeze(mask)), axis=0)
      token_acc = tf.reduce_sum(tf.cast(correct, tf.float32) * mask)
      token_acc /= tf.reduce_sum(sequence_sizes)

      # Compute Loss.
      mask = tf.cast(tf.tile(tf.expand_dims(mask, 2), (1, 1, logits.shape[2])),
                     tf.float32)
      masked_logits = logits * mask
      masked_target = tf.cast(targets, tf.float32) * mask
      logits_flat = tf.reshape(masked_logits,
                               [sq_sz_out_max * batch_size, -1])
      target_flat = tf.reshape(masked_target,
                               [sq_sz_out_max * batch_size, -1])
      xent = tf.nn.softmax_cross_entropy_with_logits(logits=logits_flat,
                                                     labels=target_flat)
      loss = tf.reduce_mean(xent)
      return loss, token_acc, solved

def aggregate_single_gradient(grad_and_vars, use_mean, check_inf_nan):
    """Calculate the average gradient for a shared variable across all towers.

  Note that this function provides a synchronization point across all towers.

  Args:
    grad_and_vars: A list or tuple of (gradient, variable) tuples. Each
      (gradient, variable) pair within the outer list represents the gradient
      of the variable calculated for a single tower, and the number of pairs
      equals the number of towers.
    use_mean: if True, mean is taken, else sum of gradients is taken.
    check_inf_nan: check grads for nans and infs.

  Returns:
    The tuple ([(average_gradient, variable),], has_nan_or_inf) where the
      gradient has been averaged across all towers. The variable is chosen from
      the first tower. The has_nan_or_inf indicates the grads has nan or inf.
  """
    grads = [g for g, _ in grad_and_vars]
    grad = tf.add_n(grads)

    if use_mean and len(grads) > 1:
        grad = tf.multiply(grad, 1.0 / len(grads))

    v = grad_and_vars[0][1]
    if check_inf_nan:
        has_nan_or_inf = tf.logical_not(tf.reduce_all(tf.is_finite(grads)))
        return (grad, v), has_nan_or_inf
    else:
        return (grad, v), None

    def test_horovod_allreduce_cpu_fused(self):
        """Test on CPU that the allreduce correctly sums 1D, 2D, 3D tensors
        with Tensor Fusion."""
        hvd.init()
        size = hvd.size()
        with self.test_session() as session:
            dtypes = [tf.int32, tf.int64, tf.float32, tf.float64]
            dims = [1, 2, 3]
            tests = []
            for dtype, dim in itertools.product(dtypes, dims):
                with tf.device("/cpu:0"):
                    tf.set_random_seed(1234)
                    tensor = tf.random_uniform(
                        [17] * dim, -100, 100, dtype=dtype)
                    summed = hvd.allreduce(tensor, average=False)
                multiplied = tensor * size
                max_difference = tf.reduce_max(tf.abs(summed - multiplied))

                # Threshold for floating point equality depends on number of
                # ranks, since we're comparing against precise multiplication.
                if size <= 3:
                    threshold = 0
                elif size < 10:
                    threshold = 1e-4
                elif size < 15:
                    threshold = 5e-4
                else:
                    break

                test = max_difference <= threshold
                tests.append(test)
            self.assertTrue(session.run(tf.reduce_all(tests)),
                            "hvd.allreduce produces incorrect results")

  def test_mpi_allgather_variable_size(self):
    """Test that the allgather correctly gathers 1D, 2D, 3D tensors,
    even if those tensors have different sizes along the first dim."""
    with self.test_session() as session:
      size = session.run(mpi.size())
      rank = session.run(mpi.rank())

      dtypes = tf.int32, tf.float32
      dims = 1, 2, 3
      for dtype, dim in itertools.product(dtypes, dims):
        # Support tests up to MPI Size of 35
        if size > 35:
          break

        tensor_sizes = [17, 32, 81, 12, 15, 23, 22] * 5
        tensor_sizes = tensor_sizes[:size]

        tensor = tf.ones([tensor_sizes[rank]] + [17] * (dim - 1),
                         dtype=dtype) * rank
        gathered = mpi.allgather(tensor)

        gathered_tensor = session.run(gathered)
        expected_size = sum(tensor_sizes)
        self.assertEqual(list(gathered_tensor.shape),
                         [expected_size] + [17] * (dim - 1))

        for i in range(size):
          rank_size = [tensor_sizes[i]] + [17] * (dim - 1)
          rank_tensor = tf.slice(gathered,
                                 [sum(tensor_sizes[:i])] + [0] * (dim - 1),
                                 rank_size)
          self.assertEqual(list(rank_tensor.shape), rank_size)
          self.assertTrue(session.run(tf.reduce_all(tf.equal(rank_tensor, i))),
                          "mpi.allgather produces incorrect gathered tensor")

def ImageSample(inputs, borderMode='repeat'):
    """
    Sample the images using the given coordinates, by bilinear interpolation.
    This was described in the paper:
    `Spatial Transformer Networks <http://arxiv.org/abs/1506.02025>`_.

    This is equivalent to `torch.nn.functional.grid_sample`,
    up to some non-trivial coordinate transformation.

    This implementation returns pixel value at pixel (1, 1) for a floating point coordinate (1.0, 1.0).
    Note that this may not be what you need.

    Args:
        inputs (list): [images, coords]. images has shape NHWC.
            coords has shape (N, H', W', 2), where each pair of the last dimension is a (y, x) real-value
            coordinate.
        borderMode: either "repeat" or "constant" (zero-filled)

    Returns:
        tf.Tensor: a tensor named ``output`` of shape (N, H', W', C).
    """
    log_deprecated("ImageSample", "Please implement it in your own code instead!", "2018-12-01")
    image, mapping = inputs
    assert image.get_shape().ndims == 4 and mapping.get_shape().ndims == 4
    input_shape = image.get_shape().as_list()[1:]
    assert None not in input_shape, \
        "Images in ImageSample layer must have fully-defined shape"
    assert borderMode in ['repeat', 'constant']

    orig_mapping = mapping
    mapping = tf.maximum(mapping, 0.0)
    lcoor = tf.floor(mapping)
    ucoor = lcoor + 1

    diff = mapping - lcoor
    neg_diff = 1.0 - diff  # bxh2xw2x2

    lcoory, lcoorx = tf.split(lcoor, 2, 3)
    ucoory, ucoorx = tf.split(ucoor, 2, 3)

    lyux = tf.concat([lcoory, ucoorx], 3)
    uylx = tf.concat([ucoory, lcoorx], 3)

    diffy, diffx = tf.split(diff, 2, 3)
    neg_diffy, neg_diffx = tf.split(neg_diff, 2, 3)

    ret = tf.add_n([sample(image, lcoor) * neg_diffx * neg_diffy,
                    sample(image, ucoor) * diffx * diffy,
                    sample(image, lyux) * neg_diffy * diffx,
                    sample(image, uylx) * diffy * neg_diffx], name='sampled')
    if borderMode == 'constant':
        max_coor = tf.constant([input_shape[0] - 1, input_shape[1] - 1], dtype=tf.float32)
        mask = tf.greater_equal(orig_mapping, 0.0)
        mask2 = tf.less_equal(orig_mapping, max_coor)
        mask = tf.logical_and(mask, mask2)  # bxh2xw2x2
        mask = tf.reduce_all(mask, [3])  # bxh2xw2 boolean
        mask = tf.expand_dims(mask, 3)
        ret = ret * tf.cast(mask, tf.float32)
    return tf.identity(ret, name='output')

 def flip_img_coord(_img, _coord):
     zeros = tf.constant([[0, 0, 0, 0, 0]]*30, tf.float32)
     img_flipped = tf.image.flip_left_right(_img)
     idx_invalid = tf.reduce_all(tf.equal(coord, 0), axis=-1)
     coord_temp = tf.concat([tf.minimum(tf.maximum(1 - _coord[:, :1], 0), 1),
                            _coord[:, 1:]], axis=-1)
     coord_flipped = tf.where(idx_invalid, zeros, coord_temp)
     return img_flipped, coord_flipped

def custom_dynamic_rnn(cell, inputs, inputs_len, initial_state=None):
    """
    Implements a dynamic rnn that can store scores in the pointer network,
    the reason why we implements this is that the raw_rnn or dynamic_rnn function in Tensorflow
    seem to require the hidden unit and memory unit has the same dimension, and we cannot
    store the scores directly in the hidden unit.
    Args:
        cell: RNN cell
        inputs: the input sequence to rnn
        inputs_len: valid length
        initial_state: initial_state of the cell
    Returns:
        outputs and state
    """
    batch_size = tf.shape(inputs)[0]
    max_time = tf.shape(inputs)[1]

    inputs_ta = tf.TensorArray(dtype=tf.float32, size=max_time)
    inputs_ta = inputs_ta.unstack(tf.transpose(inputs, [1, 0, 2]))
    emit_ta = tf.TensorArray(dtype=tf.float32, dynamic_size=True, size=0)
    t0 = tf.constant(0, dtype=tf.int32)
    if initial_state is not None:
        s0 = initial_state
    else:
        s0 = cell.zero_state(batch_size, dtype=tf.float32)
    f0 = tf.zeros([batch_size], dtype=tf.bool)

    def loop_fn(t, prev_s, emit_ta, finished):
        """
        the loop function of rnn
        """
        cur_x = inputs_ta.read(t)
        scores, cur_state = cell(cur_x, prev_s)

        # copy through
        scores = tf.where(finished, tf.zeros_like(scores), scores)

        if isinstance(cell, tc.rnn.LSTMCell):
            cur_c, cur_h = cur_state
            prev_c, prev_h = prev_s
            cur_state = tc.rnn.LSTMStateTuple(tf.where(finished, prev_c, cur_c),
                                              tf.where(finished, prev_h, cur_h))
        else:
            cur_state = tf.where(finished, prev_s, cur_state)

        emit_ta = emit_ta.write(t, scores)
        finished = tf.greater_equal(t + 1, inputs_len)
        return [t + 1, cur_state, emit_ta, finished]

    _, state, emit_ta, _ = tf.while_loop(
        cond=lambda _1, _2, _3, finished: tf.logical_not(tf.reduce_all(finished)),
        body=loop_fn,
        loop_vars=(t0, s0, emit_ta, f0),
        parallel_iterations=32,
        swap_memory=False)

    outputs = tf.transpose(emit_ta.stack(), [1, 0, 2])
    return outputs, state

    def _resample_nearest(self, inputs, sample_coords):
        # This is forward only as no gradient for tf.round

        # read input shape
        in_size = inputs.shape
        partial_shape = not in_size.is_fully_defined()
        in_spatial_size = None

        try:
            batch_size = int(in_size[0])
            n_coords = int(sample_coords.shape[0])
        except (TypeError, AssertionError, ValueError):
            tf.logging.fatal('Unknown input shape, at least batch size '
                             'and rank of the inputs are required.')
            raise

        # quantise coordinates
        spatial_coords = tf.round(sample_coords)
        if not partial_shape:
            in_spatial_size = in_size.as_list()[1:-1]
            spatial_coords = self.boundary_func(
                spatial_coords, in_spatial_size)
        spatial_coords = tf.cast(spatial_coords, COORDINATES_TYPE)

        if batch_size == n_coords:
            batch_inputs = tf.unstack(inputs)
            batch_coords = tf.unstack(spatial_coords)
            gathered_image = [tf.gather_nd(img, coord) for (img, coord) in
                              zip(batch_inputs, batch_coords)]
        elif n_coords == 1 and batch_size > 1:
            gathered_image = [tf.gather_nd(img, spatial_coords[0])
                              for img in tf.unstack(inputs)]
        else:
            raise NotImplementedError
        output = tf.stack(gathered_image, axis=0)

        if self.boundary == 'ZERO' and in_spatial_size:
            scale = 1. / (tf.constant(in_spatial_size, dtype=tf.float32) - 1.)
            mask = tf.logical_and(
                tf.reduce_all(sample_coords > 0,
                              axis=-1, keep_dims=True),
                tf.reduce_all(scale * sample_coords < 1,
                              axis=-1, keep_dims=True))
            return output * tf.to_float(mask)
        return output

def safe_sum(x, alt_value=-np.inf, name=None):
  """Elementwise adds list members, replacing non-finite results with alt_value.

  Args:
    x: Python `list` of `Tensors` to elementwise add.
    alt_value: Python scalar used to replace any elementwise sums which would
      otherwise be non-finite.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "safe_sum").

  Returns:
    safe_sum: `Tensor` representing the elementwise sum of list of `Tensor`s
      `x` or `alt_value` where sums are non-finite.

  Raises:
    TypeError: if `x` is not list-like.
    ValueError: if `x` is empty.
  """
  with tf.name_scope(name, 'safe_sum', [x, alt_value]):
    if not is_list_like(x):
      raise TypeError('Expected list input.')
    if not x:
      raise ValueError('Input should not be empty.')
    n = np.int32(len(x))
    in_shape = x[0].shape
    x = tf.stack(x, axis=-1)
    # The sum is NaN if any element is NaN or we see both +Inf and -Inf.  Thus
    # we will replace such rows with the `alt_value`. Typically the `alt_value`
    # is chosen so the `MetropolisHastings` `TransitionKernel` always rejects
    # the proposal.  rejection.
    # Regarding the following float-comparisons, recall comparing with NaN is
    # always False, i.e., we're implicitly capturing NaN and explicitly
    # capturing +/- Inf.
    is_sum_determinate = (
        tf.reduce_all(tf.is_finite(x) | (x >= 0.), axis=-1) &
        tf.reduce_all(tf.is_finite(x) | (x <= 0.), axis=-1))
    is_sum_determinate = tf.tile(
        is_sum_determinate[..., tf.newaxis],
        multiples=tf.concat([tf.ones(tf.rank(x) - 1, dtype=tf.int32), [n]],
                            axis=0))
    alt_value = np.array(alt_value, x.dtype.as_numpy_dtype)
    x = tf.where(is_sum_determinate, x, tf.fill(tf.shape(x), value=alt_value))
    x = tf.reduce_sum(x, axis=-1)
    x.set_shape(x.shape.merge_with(in_shape))
    return x

def all(x, axis=None, keepdims=False):
    '''Bitwise reduction (logical AND).

    Returns an uint8 tensor
    '''
    axis = _normalize_axis(axis, ndim(x))
    x = tf.cast(x, tf.bool)
    x = tf.reduce_all(x, reduction_indices=axis, keep_dims=keepdims)
    return tf.cast(x, tf.uint8)