jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_horovod_grouped_allgather_cpu(self):
        """Test that the grouped allgather correctly gathers 1D, 2D, 3D tensors."""
        hvd.init()
        rank = hvd.rank()
        size = hvd.size()

        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)

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

            for i in range(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 = i
                else:
                    value = 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 sinusoidal_embedding(self, n_channels, dim):
        pe = t.cast(([[p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)] for p in range(n_channels)]), t.float32)
        pe[:, 0::2] = t.sin(pe[:, 0::2])
        pe[:, 1::2] = t.cos(pe[:, 1::2])
        pe = tf.expand_dims(pe, axis=0)

        return pe
//  [python]
def sinusoidal_embedding(self, n_channels, dim):
        pe = tf.cast(([[p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)] for p in range(n_channels)]), tf.float32)
        pe[:, 0::2] = tf.sin(pe[:, 0::2])
        pe[:, 1::2] = tf.cos(pe[:, 1::2])
        pe = tf.expand_dims(pe, axis=0)

        return pe
//  [python]
def top_k_accuracy(y_true, y_pred, k=1, tf_enabled=True):
    '''
    Calculates top_k accuracy of predictions. Expects both y_true and y_pred to be one-hot encoded.
    numpy implementation is from: https://github.com/chainer/chainer/issues/606
    '''

    if tf_enabled:
        argsorted_y = argsort(y_pred)[:,-k:]
        matches = t.cast(reduce_any(transpose(argsorted_y) == argmax(y_true, axis=1, output_type=int32), axis=0), float32)
        return  reduce_mean(matches).numpy()
    else:
        argsorted_y = np.argsort(y_pred)[:,-k:]
        return np.any(argsorted_y.T == y_true.argmax(axis=1), axis=0).mean()
//  [python]
def top_k_accuracy(y_true, y_pred, k=1, tf_enabled=True):
    '''
    Calculates top_k accuracy of predictions. Expects both y_true and y_pred to be one-hot encoded.
    numpy implementation is from: https://github.com/chainer/chainer/issues/606
    '''

    if tf_enabled:
        argsorted_y = argsort(y_pred)[:,-k:]
        matches = cast(reduce_any(transpose(argsorted_y) == argmax(y_true, axis=1, output_type=int32), axis=0), float32)
        return  reduce_mean(matches).numpy()
    else:
        argsorted_y = np.argsort(y_pred)[:,-k:]
        return np.any(argsorted_y.T == y_true.argmax(axis=1), axis=0).mean()
//  [python]
def test_horovod_grouped_allgather_gpu(self):
        """Test that the grouped allgather correctly gathers 1D, 2D, 3D tensors."""
        # 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")

        hvd.init()
        rank = hvd.rank()
        local_rank = hvd.local_rank()
        size = hvd.size()

        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("/gpu:%d" % local_rank):
                gathered = hvd.grouped_allgather(tensors)

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

            for i in range(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 = i
                else:
                    value = 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 mask_invalid(coords, pad_h=0, pad_w=0):
  """Mask coordinates outside of the image.

  Valid = 1, invalid = 0.

  Args:
    coords: a 4D float tensor of image coordinates.
    pad_h: int, the amount of padding applied to the top of the image
    pad_w: int, the amount of padding applied to the left of the image

  Returns:
    The mask showing which coordinates are valid.
  """
  pad_h = float(pad_h)
  pad_w = float(pad_w)
  coords_rank = len(coords.shape)
  if coords_rank != 4:
    raise NotImplementedError()
  max_height = float(coords.shape[-3] - 1)
  max_width = float(coords.shape[-2] - 1)
  mask = tf.logical_and(
      tf.logical_and(coords[:, :, :, 0] >= pad_h,
                     coords[:, :, :, 0] <= max_height),
      tf.logical_and(coords[:, :, :, 1] >= pad_w,
                     coords[:, :, :, 1] <= max_width))
  mask = tf.to_float(mask)[:, :, :, None]
  return mask
//  [python]
def mask_invalid(coords, pad_h=0, pad_w=0):
  """Mask coordinates outside of the image.

  Valid = 1, invalid = 0.

  Args:
    coords: a 4D float tensor of image coordinates.
    pad_h: int, the amount of padding applied to the top of the image
    pad_w: int, the amount of padding applied to the left of the image

  Returns:
    The mask showing which coordinates are valid.
  """
  pad_h = float(pad_h)
  pad_w = float(pad_w)
  coords_rank = len(coords.shape)
  if coords_rank != 4:
    raise NotImplementedError()
  max_height = float(coords.shape[-3] - 1)
  max_width = float(coords.shape[-2] - 1)
  mask = tf.logical_and(
      tf.logical_and(coords[:, :, :, 0] >= pad_h,
                     coords[:, :, :, 0] <= max_height),
      tf.logical_and(coords[:, :, :, 1] >= pad_w,
                     coords[:, :, :, 1] <= max_width))
  mask = tf.cast(mask, dtype=tf.float32)[:, :, :, None]
  return mask
//  [python]
def resample(source, coords):
  """Resample the source image at the passed coordinates.

  Args:
    source: tf.tensor, batch of images to be resampled.
    coords: tf.tensor, batch of coordinates in the image. Coordinates should
      be between 0 and size - 1. Coordinates outside of this range are handled
      by interpolating with a background image filled with zeros in the same
      way that SAME size convolution works.

  Returns:
    The resampled image.
  """

  # Wrap this function because it uses a different order of height/width dims.
  orig_source_dtype = source.dtype
  if source.dtype != tf.float32:
    source = t.cast(source, tf.float32)
  if coords.dtype != tf.float32:
    coords = t.cast(coords, tf.float32)
  coords_rank = len(coords.shape)
  if coords_rank == 4:
    output = resampler(source, coords[:, :, :, ::-1])
    if orig_source_dtype != source.dtype:
      return t.cast(output, orig_source_dtype)
    return output
  else:
    raise NotImplementedError()
//  [python]
def resample(source, coords):
  """Resample the source image at the passed coordinates.

  Args:
    source: tf.tensor, batch of images to be resampled.
    coords: tf.tensor, batch of coordinates in the image. Coordinates should
      be between 0 and size - 1. Coordinates outside of this range are handled
      by interpolating with a background image filled with zeros in the same
      way that SAME size convolution works.

  Returns:
    The resampled image.
  """

  # Wrap this function because it uses a different order of height/width dims.
  orig_source_dtype = source.dtype
  if source.dtype != tf.float32:
    source = tf.cast(source, tf.float32)
  if coords.dtype != tf.float32:
    coords = tf.cast(coords, tf.float32)
  coords_rank = len(coords.shape)
  if coords_rank == 4:
    output = resampler(source, coords[:, :, :, ::-1])
    if orig_source_dtype != source.dtype:
      return tf.cast(output, orig_source_dtype)
    return output
  else:
    raise NotImplementedError()
//  [python]
def filterbank(self, gx, gy, sigma2, delta):
        # 1 x N, look like [[0,1,2,3,4]]
        grid_i = tf.reshape(tf.to_float(tf.range(self.attention_n)), [1, -1])
        # centers for the individual patches
        mu_x = gx + (grid_i - self.attention_n/2 - 0.5) * delta
        mu_y = gy + (grid_i - self.attention_n/2 - 0.5) * delta
        mu_x = tf.reshape(mu_x, [-1, self.attention_n, 1])
        mu_y = tf.reshape(mu_y, [-1, self.attention_n, 1])
        # 1 x 1 x imgsize, looks like [[[0,1,2,3,4,...,27]]]
        im_a = tf.reshape(tf.to_float(tf.range(self.img_w)), [1, 1, -1])
        im_b = tf.reshape(tf.to_float(tf.range(self.img_h)), [1, 1, -1])
        # list of gaussian curves for x and y
        sigma2 = tf.reshape(sigma2, [-1, 1, 1])
        Fx = tf.exp(-tf.square(im_a - mu_x) / (2*sigma2))
        Fy = tf.exp(-tf.square(im_b - mu_y) / (2*sigma2))
        # normalize so area-under-curve = 1
        Fx = Fx / tf.maximum(tf.reduce_sum(Fx, 2, keepdims=True), 1e-8)
        Fy = Fy / tf.maximum(tf.reduce_sum(Fy, 2, keepdims=True), 1e-8)
        return Fx, Fy
//  [python]
def filterbank(self, gx, gy, sigma2, delta):
        # 1 x N, look like [[0,1,2,3,4]]
        grid_i = tf.reshape(tf.cast(tf.range(self.attention_n), self.dtype), [1, -1])
        # centers for the individual patches
        mu_x = gx + (grid_i - self.attention_n/2 - 0.5) * delta
        mu_y = gy + (grid_i - self.attention_n/2 - 0.5) * delta
        mu_x = tf.reshape(mu_x, [-1, self.attention_n, 1])
        mu_y = tf.reshape(mu_y, [-1, self.attention_n, 1])
        # 1 x 1 x imgsize, looks like [[[0,1,2,3,4,...,27]]]
        im_a = tf.reshape(tf.range(self.img_w, dtype=self.dtype), [1, 1, -1])
        im_b = tf.reshape(tf.range(self.img_h, dtype=self.dtype), [1, 1, -1])
        # list of gaussian curves for x and y
        sigma2 = tf.reshape(sigma2, [-1, 1, 1])
        Fx = tf.exp(-tf.square(im_a - mu_x) / (2*sigma2))
        Fy = tf.exp(-tf.square(im_b - mu_y) / (2*sigma2))
        # normalize so area-under-curve = 1
        Fx = Fx / tf.maximum(tf.reduce_sum(Fx, 2, keepdims=True), 1e-8)
        Fy = Fy / tf.maximum(tf.reduce_sum(Fy, 2, keepdims=True), 1e-8)
        return Fx, Fy
//  [python]
def compute_occlusions_brox(forward_flow,
                            backward_flow):
  """Compute an occlusion mask based on a forward-backward check.

  Args:
    forward_flow: Forward flow field of shape [batch, height, width, 2].
    backward_flow: Backward flow field of shape [batch, height, width, 2].

  Returns:
    Occlusion mask of shape [batch, height, width, 1], where 1 are occluded
    locations and 0 are non-occluded.
  """
  # Resampled backward flow at forward flow locations.
  warp = flow_to_warp(forward_flow)
  backward_flow_resampled = resample(backward_flow, warp)

  # Compute occlusions based on forward-backward consistency.
  fb_sq_diff = tf.reduce_sum(
      (forward_flow + backward_flow_resampled)**2, axis=-1, keepdims=True)
  fb_sum_sq = tf.reduce_sum(
      forward_flow**2 + backward_flow_resampled**2, axis=-1, keepdims=True)
  return t.cast(fb_sq_diff > 0.01 * fb_sum_sq + 0.5, tf.float32)
//  [python]
def compute_occlusions_brox(forward_flow,
                            backward_flow):
  """Compute an occlusion mask based on a forward-backward check.

  Args:
    forward_flow: Forward flow field of shape [batch, height, width, 2].
    backward_flow: Backward flow field of shape [batch, height, width, 2].

  Returns:
    Occlusion mask of shape [batch, height, width, 1], where 1 are occluded
    locations and 0 are non-occluded.
  """
  # Resampled backward flow at forward flow locations.
  warp = flow_to_warp(forward_flow)
  backward_flow_resampled = resample(backward_flow, warp)

  # Compute occlusions based on forward-backward consistency.
  fb_sq_diff = tf.reduce_sum(
      (forward_flow + backward_flow_resampled)**2, axis=-1, keepdims=True)
  fb_sum_sq = tf.reduce_sum(
      forward_flow**2 + backward_flow_resampled**2, axis=-1, keepdims=True)
  return tf.cast(fb_sq_diff > 0.01 * fb_sum_sq + 0.5, tf.float32)
//  [python]
def compute_occlusions_wang(backward_flow, downsampling_factor,
                            threshold):
  """Compute occlusion mask based on a rangemap.

  Args:
    backward_flow: Backward flow field of shape [batch, height, width, 2].
    downsampling_factor: Downsampling factor used for the range map computation.
    threshold: Indicates if thresholding should be used

  Returns:
    Occlusion mask of shape [batch, height, width, 1], where 1 are occluded
    locations and 0 are non-occluded.
  """
  range_map = compute_range_map(
      backward_flow,
      downsampling_factor=downsampling_factor,
      reduce_downsampling_bias=False,
      resize_output=False)
  if threshold:
    return 1.0 - t.cast(range_map < 0.75, t.float32)
  else:
    return 1.0 - tf.clip_by_value(range_map, 0.0, 1.0)
//  [python]
def compute_occlusions_wang(backward_flow, downsampling_factor,
                            threshold):
  """Compute occlusion mask based on a rangemap.

  Args:
    backward_flow: Backward flow field of shape [batch, height, width, 2].
    downsampling_factor: Downsampling factor used for the range map computation.
    threshold: Indicates if thresholding should be used

  Returns:
    Occlusion mask of shape [batch, height, width, 1], where 1 are occluded
    locations and 0 are non-occluded.
  """
  range_map = compute_range_map(
      backward_flow,
      downsampling_factor=downsampling_factor,
      reduce_downsampling_bias=False,
      resize_output=False)
  if threshold:
    return 1.0 - tf.cast(range_map < 0.75, tf.float32)
  else:
    return 1.0 - tf.clip_by_value(range_map, 0.0, 1.0)
//  [python]
def dot_product_attention(q, k, v, bias, dropout_rate=0.0):
  """Dot-product attention.
  Args:
    q: Tensor with shape [..., length_q, depth_k].
    k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
      match with q.
    v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
      match with q.
    bias: bias Tensor (see attention_bias())
    dropout_rate: a float.
  Returns:
    Tensor with shape [..., length_q, depth_v].
  """
  logits = tf.matmul(q, k, transpose_b=True)  # [..., length_q, length_kv]
  logits = tf.multiply(logits, 1.0 / math.sqrt(float(get_shape_list(q)[-1])))
  if bias is not None:
    # `attention_mask` = [B, T]
    from_shape = get_shape_list(q)
    if len(from_shape) == 4:
      broadcast_ones = tf.ones([from_shape[0], 1, from_shape[2], 1], tf.float32)
    elif len(from_shape) == 5:
      # from_shape = [B, N, Block_num, block_size, depth]#
      broadcast_ones = tf.ones([from_shape[0], 1, from_shape[2], from_shape[3],
                                1], tf.float32)

    bias = tf.matmul(broadcast_ones,
                     tf.to_float(bias), transpose_b=True)

    # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
    # masked positions, this operation will create a tensor which is 0.0 for
    # positions we want to attend and -10000.0 for masked positions.
    adder = (1.0 - bias) * -10000.0

    # Since we are adding it to the raw scores before the softmax, this is
    # effectively the same as removing these entirely.
    logits += adder
  else:
    adder = 0.0

  attention_probs = tf.nn.softmax(logits, name="attention_probs")
  attention_probs = dropout(attention_probs, dropout_rate)
  return tf.matmul(attention_probs, v)
//  [python]
def dot_product_attention(q, k, v, bias, dropout_rate=0.0):
  """Dot-product attention.
  Args:
    q: Tensor with shape [..., length_q, depth_k].
    k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
      match with q.
    v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
      match with q.
    bias: bias Tensor (see attention_bias())
    dropout_rate: a float.
  Returns:
    Tensor with shape [..., length_q, depth_v].
  """
  logits = tf.matmul(q, k, transpose_b=True)  # [..., length_q, length_kv]
  logits = tf.multiply(logits, 1.0 / math.sqrt(float(get_shape_list(q)[-1])))
  if bias is not None:
    # `attention_mask` = [B, T]
    from_shape = get_shape_list(q)
    if len(from_shape) == 4:
      broadcast_ones = tf.ones([from_shape[0], 1, from_shape[2], 1], tf.float32)
    elif len(from_shape) == 5:
      # from_shape = [B, N, Block_num, block_size, depth]#
      broadcast_ones = tf.ones([from_shape[0], 1, from_shape[2], from_shape[3],
                                1], tf.float32)

    bias = tf.compat.v1.matmul(broadcast_ones,
                     tf.cast(bias, tf.float32), transpose_b=True)

    # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
    # masked positions, this operation will create a tensor which is 0.0 for
    # positions we want to attend and -10000.0 for masked positions.
    adder = (1.0 - bias) * -10000.0

    # Since we are adding it to the raw scores before the softmax, this is
    # effectively the same as removing these entirely.
    logits += adder
  else:
    adder = 0.0

  attention_probs = tf.nn.softmax(logits, name="attention_probs")
  attention_probs = dropout(attention_probs, dropout_rate)
  return tf.matmul(attention_probs, v)
//  [python]
def assert_shape(tensor, shape):
    assert(len(tensor.get_shape()) == len(shape))
    tensor_shape = tf.shape(tensor)
    for i, s_i in enumerate(shape):
        tf.assert_equal(tensor_shape[i], tf.to_int32(s_i))
//  [python]
def assert_shape(tensor, shape):
    assert(len(tensor.get_shape()) == len(shape))
    tensor_shape = tf.shape(tensor)
    for i, s_i in enumerate(shape):
        tf.assert_equal(tensor_shape[i], tf.cast(s_i, tf.int32))
//  [python]
def assert_equal_approx(tensor, value, eps=1e-5, name=None):
    return tf.assert_equal(tf.cast(tf.abs(tensor - value) < 1e-5, tf.int32), 1, name=name)
//  [python]
def supervised_sequence_loss(ground_truth_flow,
                             ground_truth_valid,
                             predicted_flows,
                             weights,
                             max_flow_threshold=400,
                             max_flow_norm='l2',
                             loss_decay=.8):
  """Computes a supervised sequence loss."""
  del weights
  losses = {}
  def fn(accum, predicted_flow):
    height = tf.shape(ground_truth_flow)[-3]
    width = tf.shape(ground_truth_flow)[-2]
    upsampled_flow = resize(predicted_flow, height, width, is_flow=True)

    if max_flow_norm == 'l1':
      # Version used in earlier RAFT implementation with a max_flow_threshold
      # of 1000 and later 500.
      max_flow_mask = tf.to_float(
          tf.norm(ground_truth_flow, ord=1, axis=-1, keepdims=True) <
          max_flow_threshold)
    elif max_flow_norm == 'l2':
      max_flow_mask = tf.to_float(
          tf.norm(ground_truth_flow, ord='euclidean', axis=-1, keepdims=True) <
          max_flow_threshold)
    else:
      raise ValueError('The value of max_flow_norm within the supervised '
                       'sequence loss must be either l1 or l2.')

    loss = tf.math.abs(ground_truth_flow - upsampled_flow)
    loss = tf.reduce_mean(loss * max_flow_mask * ground_truth_valid)
    return accum * loss_decay + loss

  loss_array = tf.scan(
      fn,
      predicted_flows,
      initializer=tf.convert_to_tensor([0], dtype=tf.float32),
      parallel_iterations=1)

  final_accum_loss = loss_array[-1, 0]
  losses['supervised_sequence_loss'] = final_accum_loss
  return losses
//  [python]
def supervised_sequence_loss(ground_truth_flow,
                             ground_truth_valid,
                             predicted_flows,
                             weights,
                             max_flow_threshold=400,
                             max_flow_norm='l2',
                             loss_decay=.8):
  """Computes a supervised sequence loss."""
  del weights
  losses = {}
  def fn(accum, predicted_flow):
    height = tf.shape(ground_truth_flow)[-3]
    width = tf.shape(ground_truth_flow)[-2]
    upsampled_flow = resize(predicted_flow, height, width, is_flow=True)

    if max_flow_norm == 'l1':
      # Version used in earlier RAFT implementation with a max_flow_threshold
      # of 1000 and later 500.
      max_flow_mask = tf.cast(
          tf.norm(ground_truth_flow, ord=1, axis=-1, keepdims=True) <
          max_flow_threshold, tf.float32)
    elif max_flow_norm == 'l2':
      max_flow_mask = tf.cast(
          tf.norm(ground_truth_flow, ord='euclidean', axis=-1, keepdims=True) <
          max_flow_threshold, tf.float32)
    else:
      raise ValueError('The value of max_flow_norm within the supervised '
                       'sequence loss must be either l1 or l2.')

    loss = tf.math.abs(ground_truth_flow - upsampled_flow)
    loss = tf.reduce_mean(loss * max_flow_mask * ground_truth_valid)
    return accum * loss_decay + loss

  loss_array = tf.scan(
      fn,
      predicted_flows,
      initializer=tf.convert_to_tensor([0], dtype=tf.float32),
      parallel_iterations=1)

  final_accum_loss = loss_array[-1, 0]
  losses['supervised_sequence_loss'] = final_accum_loss
  return losses
//  [python]
def _mask_by_length(t, length):
  """Mask t, 3-D [batch, time, dim], by length, 1-D [batch,]."""
  maxlen = t.get_shape().as_list()[1]

  # Subtract 1 from length to prevent the perturbation from going on 'eos'
  mask = tf.sequence_mask(length - 1, maxlen=maxlen)
  mask = tf.expand_dims(t.astype(tf.float32), -1)
  # shape(mask) = (batch, num_timesteps, 1)
  return t * mask
//  [python]
def _mask_by_length(t, length):
  """Mask t, 3-D [batch, time, dim], by length, 1-D [batch,]."""
  maxlen = t.get_shape().as_list()[1]

  # Subtract 1 from length to prevent the perturbation from going on 'eos'
  mask = tf.sequence_mask(length - 1, maxlen=maxlen)
  mask = tf.expand_dims(tf.cast(mask, tf.float32), -1)
  # shape(mask) = (batch, num_timesteps, 1)
  return t * mask
//  [python]
def predict(self, image_arrays):
    input_details = self.model.get_input_details()
    output_details = self.model.get_output_details()
    input_detail = input_details[0]
    if input_detail['quantization'] != (DEFAULT_SCALE, DEFAULT_ZERO_POINT):
      scale, zero_point = input_detail['quantization']
      image_arrays = image_arrays / scale + zero_point
      image_arrays = tf.cast(image_arrays, dtype=input_detail['dtype'])

    signature = list(self.model.get_signature_list().keys())[0]
    infer_fn = self.model.get_signature_runner(signature)
    outputs = infer_fn(images=image_arrays)

    def get_output(output_detail, output_tensor):
      if output_detail['quantization'] != (DEFAULT_SCALE, DEFAULT_ZERO_POINT):
        # Dequantize the output
        scale, zero_point = output_detail['quantization']
        output_tensor = output_tensor.astype(np.float32)
        output_tensor = (output_tensor - zero_point) * scale
      return output_tensor

    outputs = [
        get_output(output_detail, output)
        for output_detail, output in zip(output_details, outputs.values())
    ]
    if self.only_network:
      outputs = tuple(outputs)
    return outputs
//  [python]
def spatial_soft_argmax(features, dtype=tf.float32):
    """
    features shape is [N, H, W, C]
    """
    N = tf.shape(features)[0]
    val_shape = features.get_shape()
    H, W, C = val_shape[1].value, val_shape[2].value, val_shape[3].value
    features = tf.reshape(
        tf.transpose(features, [0, 3, 1, 2]),
        [-1, H * W])
    softmax = tf.nn.softmax(features)
    spatial_softmax = tf.transpose(tf.reshape(softmax, [N, C, H, W]), [0, 2, 3, 1])
    spatial_softmax_pos = tf.expand_dims(spatial_softmax, -1)
    # TODO shape [H, W, 1, 2]
    # TODO H or W is 1
    assert(H != 1 and W != 1)
    delta_h = 2. / t.cast(H - 1, dtype)
    delta_w = 2. / t.cast(W - 1, dtype)
    ran_h = tf.tile(tf.expand_dims(tf.range(-1., 1. + delta_h, delta_h, dtype=dtype), 1), [1, W])
    ran_w = tf.tile(tf.expand_dims(tf.range(-1., 1 + delta_w, delta_w, dtype=dtype), 0), [H, 1])
    image_pos = tf.expand_dims(tf.stack([ran_h, ran_w], 2), 2)
    spatial_soft_amax = tf.reduce_sum(spatial_softmax_pos * image_pos, axis=[1, 2])
    shaped_ssamax = tf.reshape(spatial_soft_amax, [N, C * 2])
    return shaped_ssamax
//  [python]
def spatial_soft_argmax(features, dtype=tf.float32):
    """
    features shape is [N, H, W, C]
    """
    N = tf.shape(features)[0]
    val_shape = features.get_shape()
    H, W, C = val_shape[1].value, val_shape[2].value, val_shape[3].value
    features = tf.reshape(
        tf.transpose(features, [0, 3, 1, 2]),
        [-1, H * W])
    softmax = tf.nn.softmax(features)
    spatial_softmax = tf.transpose(tf.reshape(softmax, [N, C, H, W]), [0, 2, 3, 1])
    spatial_softmax_pos = tf.expand_dims(spatial_softmax, -1)
    # TODO shape [H, W, 1, 2]
    # TODO H or W is 1
    assert(H != 1 and W != 1)
    delta_h = 2. / tf.cast(H - 1, dtype)
    delta_w = 2. / tf.cast(W - 1, dtype)
    ran_h = tf.tile(tf.expand_dims(tf.range(-1., 1. + delta_h, delta_h, dtype=dtype), 1), [1, W])
    ran_w = tf.tile(tf.expand_dims(tf.range(-1., 1 + delta_w, delta_w, dtype=dtype), 0), [H, 1])
    image_pos = tf.expand_dims(tf.stack([ran_h, ran_w], 2), 2)
    spatial_soft_amax = tf.reduce_sum(spatial_softmax_pos * image_pos, axis=[1, 2])
    shaped_ssamax = tf.reshape(spatial_soft_amax, [N, C * 2])
    return shaped_ssamax
//  [python]
def minibatch(self, dataset, subset):
    with tf.compat.v1.name_scope('batch_processing'):
      images = [[] for i in range(self.device_count)]
      labels = [[] for i in range(self.device_count)]
      record_input = data_flow_ops.RecordInput(
          file_pattern=dataset.tf_record_pattern(subset),
          seed=randint(0, 9000),
          parallelism=64,
          buffer_size=10000,
          batch_size=self.batch_size,
          name='record_input')
      records = record_input.get_yield_op()
      records = tf.split(records, self.batch_size, 0)
      records = [tf.reshape(record, []) for record in records]
      for i in xrange(self.batch_size):
        value = records[i]
        image_buffer, label_index, bbox, _ = parse_example_proto(value)
        image = self.preprocess(image_buffer, bbox, i % 4)
        device_index = i % self.device_count
        images[device_index].append(image)
        labels[device_index].append(label_index)
      label_index_batch = [None] * self.device_count
      for device_index in xrange(self.device_count):
        images[device_index] = tf.parallel_stack(images[device_index])
        label_index_batch[device_index] = tf.concat(labels[device_index], 0)

        # dynamic_pad=True) # HACK TESTING dynamic_pad=True
        images[device_index] = tf.to_float(images[device_index])
        depth = 3
        images[device_index] = tf.reshape(
            images[device_index],
            shape=[self.batch_size_per_device, self.height, self.width, depth])
        label_index_batch[device_index] = tf.reshape(
            label_index_batch[device_index], [self.batch_size_per_device])
        # Display the training images in the visualizer.
        # tf.summary.image('images', images)

      return images, label_index_batch, records
//  [python]
def minibatch(self, dataset, subset):
    with tf.compat.v1.name_scope('batch_processing'):
      images = [[] for i in range(self.device_count)]
      labels = [[] for i in range(self.device_count)]
      record_input = data_flow_ops.RecordInput(
          file_pattern=dataset.tf_record_pattern(subset),
          seed=randint(0, 9000),
          parallelism=64,
          buffer_size=10000,
          batch_size=self.batch_size,
          name='record_input')
      records = record_input.get_yield_op()
      records = tf.split(records, self.batch_size, 0)
      records = [tf.reshape(record, []) for record in records]
      for i in xrange(self.batch_size):
        value = records[i]
        image_buffer, label_index, bbox, _ = parse_example_proto(value)
        image = self.preprocess(image_buffer, bbox, i % 4)
        device_index = i % self.device_count
        images[device_index].append(image)
        labels[device_index].append(label_index)
      label_index_batch = [None] * self.device_count
      for device_index in xrange(self.device_count):
        images[device_index] = tf.parallel_stack(images[device_index])
        label_index_batch[device_index] = tf.concat(labels[device_index], 0)

        # dynamic_pad=True) # HACK TESTING dynamic_pad=True
        images[device_index] = tf.cast(images[device_index], self.dtype)
        depth = 3
        images[device_index] = tf.reshape(
            images[device_index],
            shape=[self.batch_size_per_device, self.height, self.width, depth])
        label_index_batch[device_index] = tf.reshape(
            label_index_batch[device_index], [self.batch_size_per_device])
        # Display the training images in the visualizer.
        # tf.summary.image('images', images)

      return images, label_index_batch, records
//  [python]
def _build_inputs(self, image):
    """Builds classification model inputs for serving."""
    # Center crops and resizes image.
    if isinstance(image, tf.RaggedTensor):
      image = image.to_tensor()
    image = t.cast(image, dtype=tf.float32)

    # For these input types, decode_image already performs cropping.
    if not (
        self._input_type in ['tf_example', 'image_bytes']
        and len(self._input_image_size) == 2):
      image = self._crop_and_resize(image)

    # Normalizes image with mean and std pixel values.
    image = preprocess_ops.normalize_image(
        image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB)
    return image
//  [python]
def _build_inputs(self, image):
    """Builds classification model inputs for serving."""
    # Center crops and resizes image.
    if isinstance(image, tf.RaggedTensor):
      image = image.to_tensor()
    image = tf.cast(image, dtype=tf.float32)

    # For these input types, decode_image already performs cropping.
    if not (
        self._input_type in ['tf_example', 'image_bytes']
        and len(self._input_image_size) == 2):
      image = self._crop_and_resize(image)

    # Normalizes image with mean and std pixel values.
    image = preprocess_ops.normalize_image(
        image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB)
    return image
//  [python]
def read_and_preprocess(example_data):
    parsed = tf.parse_single_example(example_data, {
      'image/encoded': tf.FixedLenFeature((), tf.string, ''),
      'image/class/label': tf.FixedLenFeature([], tf.int64, 1),
    })
    image_bytes = tf.reshape(parsed['image/encoded'], shape=[])
    label = tf.to_int32(
      tf.reshape(parsed['image/class/label'], shape=[])) - 1

    # end up with pixel values that are in the -1, 1 range
    image = tf.image.decode_jpeg(image_bytes, channels=NUM_CHANNELS)
    image = tf.image.convert_image_dtype(image, dtype=tf.float32) # 0-1
    image = tf.expand_dims(image, 0) # resize_bilinear needs batches

    image = tf.image.resize_bilinear(
      image, [HEIGHT + 10, WIDTH + 10], align_corners=False)
    image = tf.squeeze(image)  # remove batch dimension
    image = tf.random_crop(image, [HEIGHT, WIDTH, NUM_CHANNELS])
    image = tf.image.random_flip_left_right(image)
    image = tf.image.random_brightness(image, max_delta=63.0 / 255.0)
    image = tf.image.random_contrast(image, lower=0.2, upper=1.8)

        
    #pixel values are in range [0,1], convert to [-1,1]
    image = tf.subtract(image, 0.5)
    image = tf.multiply(image, 2.0)
    #return {'image':image}, label
    return image, label
//  [python]
def read_and_preprocess(example_data):
    parsed = tf.parse_single_example(example_data, {
      'image/encoded': tf.FixedLenFeature((), tf.string, ''),
      'image/class/label': tf.FixedLenFeature([], tf.int64, 1),
    })
    image_bytes = tf.reshape(parsed['image/encoded'], shape=[])
    label = tf.cast(
      tf.reshape(parsed['image/class/label'], shape=[]), dtype=tf.int32) - 1

    # end up with pixel values that are in the -1, 1 range
    image = tf.image.decode_jpeg(image_bytes, channels=NUM_CHANNELS)
    image = tf.image.convert_image_dtype(image, dtype=tf.float32) # 0-1
    image = tf.expand_dims(image, 0) # resize_bilinear needs batches

    image = tf.image.resize_bilinear(
      image, [HEIGHT + 10, WIDTH + 10], align_corners=False)
    image = tf.squeeze(image)  # remove batch dimension
    image = tf.random_crop(image, [HEIGHT, WIDTH, NUM_CHANNELS])
    image = tf.image.random_flip_left_right(image)
    image = tf.image.random_brightness(image, max_delta=63.0 / 255.0)
    image = tf.image.random_contrast(image, lower=0.2, upper=1.8)

        
    #pixel values are in range [0,1], convert to [-1,1]
    image = tf.subtract(image, 0.5)
    image = tf.multiply(image, 2.0)
    #return {'image':image}, label
    return image, label
//  [python]
def image_classifier(features, labels, mode, params):
  image = features
  if isinstance(features, dict):
    image = features['image']

  ylogits, nclasses = cnn_model(image, mode, params)

  probabilities = tf.nn.softmax(ylogits)
  class_int = tf.to_int32(tf.argmax(probabilities, 1))
  class_str = tf.gather(LIST_OF_LABELS, class_int)
  
  if mode == tf.estimator.ModeKeys.TRAIN or mode == tf.estimator.ModeKeys.EVAL:
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(
        logits=ylogits, labels=tf.one_hot(labels, nclasses)))

    def metric_fn(class_int, labels):
      return {'accuracy': tf.metrics.accuracy(class_int, labels)}
    evalmetrics = (metric_fn, [class_int, labels])

    if mode == tf.estimator.ModeKeys.TRAIN:
      # this is needed for batch normalization, but has no effect otherwise
      update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
      optimizer = tf.train.AdamOptimizer(learning_rate=params['learning_rate'])
      if params['use_tpu']:
        optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer) # TPU change 1
      with tf.control_dependencies(update_ops):
         train_op = optimizer.minimize(loss, tf.train.get_global_step())
    else:
      train_op = None
  else:
    loss = None
    train_op = None
    evalmetrics = None

  return tf.contrib.tpu.TPUEstimatorSpec(  # TPU change 2
        mode=mode,
        predictions={"probabilities": probabilities, 
                     "classid": class_int, "class": class_str},
        loss=loss,
        train_op=train_op,
        eval_metrics=evalmetrics,
        export_outputs={'classes': tf.estimator.export.PredictOutput(
            {"probabilities": probabilities, "classid": class_int, 
             "class": class_str})}
    )
//  [python]
def image_classifier(features, labels, mode, params):
  image = features
  if isinstance(features, dict):
    image = features['image']

  ylogits, nclasses = cnn_model(image, mode, params)

  probabilities = tf.nn.softmax(ylogits)
  class_int = tf.cast(tf.argmax(probabilities, 1), tf.int32)
  class_str = tf.gather(LIST_OF_LABELS, class_int)
  
  if mode == tf.estimator.ModeKeys.TRAIN or mode == tf.estimator.ModeKeys.EVAL:
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(
        logits=ylogits, labels=tf.one_hot(labels, nclasses)))

    def metric_fn(class_int, labels):
      return {'accuracy': tf.metrics.accuracy(class_int, labels)}
    evalmetrics = (metric_fn, [class_int, labels])

    if mode == tf.estimator.ModeKeys.TRAIN:
      # this is needed for batch normalization, but has no effect otherwise
      update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
      optimizer = tf.train.AdamOptimizer(learning_rate=params['learning_rate'])
      if params['use_tpu']:
        optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer) # TPU change 1
      with tf.control_dependencies(update_ops):
         train_op = optimizer.minimize(loss, tf.train.get_global_step())
    else:
      train_op = None
  else:
    loss = None
    train_op = None
    evalmetrics = None

  return tf.contrib.tpu.TPUEstimatorSpec(  # TPU change 2
        mode=mode,
        predictions={"probabilities": probabilities, 
                     "classid": class_int, "class": class_str},
        loss=loss,
        train_op=train_op,
        eval_metrics=evalmetrics,
        export_outputs={'classes': tf.estimator.export.PredictOutput(
            {"probabilities": probabilities, "classid": class_int, 
             "class": class_str})}
    )
//  [python]
def _argmax_sample(logits, ignore_ids=None):
    """
    Does top-k sampling. if ignore_ids is on, then we will zero out those logits.
    :param logits: [batch_size, vocab_size] tensor
    :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
                        like padding maybe
    :param p: topp threshold to use, either a float or a [batch_size] vector
    :return: [batch_size, num_samples] samples

    """
    with tf.variable_scope('argmax_sample'):
        batch_size, vocab_size = get_shape_list(logits, expected_rank=2)

        probs = tf.nn.softmax(logits if ignore_ids is None else logits - t.cast(ignore_ids[None], tf.float32) * 1e10,
                              axis=-1)
        # [batch_size, vocab_perm]
        sample = tf.argmax(probs, 1, output_type=tf.int32)[:, None]
    return {
        'probs': probs,
        'sample': sample,
    }
//  [python]
def _argmax_sample(logits, ignore_ids=None):
    """
    Does top-k sampling. if ignore_ids is on, then we will zero out those logits.
    :param logits: [batch_size, vocab_size] tensor
    :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
                        like padding maybe
    :param p: topp threshold to use, either a float or a [batch_size] vector
    :return: [batch_size, num_samples] samples

    """
    with tf.variable_scope('argmax_sample'):
        batch_size, vocab_size = get_shape_list(logits, expected_rank=2)

        probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
                              axis=-1)
        # [batch_size, vocab_perm]
        sample = tf.argmax(probs, 1, output_type=tf.int32)[:, None]
    return {
        'probs': probs,
        'sample': sample,
    }
//  [python]
def _decode_image(self, encoded_image_bytes: str) -> tf.Tensor:
    """Decodes an image bytes to an image tensor.

    Use `tf.image.decode_image` to decode an image if input is expected to be 2D
    image; otherwise use `tf.io.decode_raw` to convert the raw bytes to tensor
    and reshape it to desire shape.

    Args:
      encoded_image_bytes: An encoded image string to be decoded.

    Returns:
      A decoded image tensor.
    """
    if len(self._input_image_size) == 2:
      # Decode an image if 2D input is expected.
      image_tensor = tf.image.decode_image(
          encoded_image_bytes, channels=self._num_channels
      )
      image_tensor.set_shape((None, None, self._num_channels))
      # Crop the image inside the same loop as decoding an image
      # if there could be several images of different sizes in the batch.
      image_tensor = tf.to_float(image_tensor)
      image_tensor = self._crop_and_resize(image_tensor)
      image_tensor = tf.to_int64(image_tensor)
      return image_tensor
    else:
      # Convert raw bytes into a tensor and reshape it, if not 2D input.
      image_tensor = tf.io.decode_raw(encoded_image_bytes, out_type=tf.uint8)
      image_tensor = tf.reshape(
          image_tensor, self._input_image_size + [self._num_channels]
      )
    return image_tensor
//  [python]
def _decode_image(self, encoded_image_bytes: str) -> tf.Tensor:
    """Decodes an image bytes to an image tensor.

    Use `tf.image.decode_image` to decode an image if input is expected to be 2D
    image; otherwise use `tf.io.decode_raw` to convert the raw bytes to tensor
    and reshape it to desire shape.

    Args:
      encoded_image_bytes: An encoded image string to be decoded.

    Returns:
      A decoded image tensor.
    """
    if len(self._input_image_size) == 2:
      # Decode an image if 2D input is expected.
      image_tensor = tf.image.decode_image(
          encoded_image_bytes, channels=self._num_channels
      )
      image_tensor.set_shape((None, None, self._num_channels))
      # Crop the image inside the same loop as decoding an image
      # if there could be several images of different sizes in the batch.
      image_tensor = tf.cast(image_tensor, dtype=tf.float32)
      image_tensor = self._crop_and_resize(image_tensor)
      image_tensor = tf.cast(image_tensor, tf.uint8)
      return image_tensor
    else:
      # Convert raw bytes into a tensor and reshape it, if not 2D input.
      image_tensor = tf.io.decode_raw(encoded_image_bytes, out_type=tf.uint8)
      image_tensor = tf.reshape(
          image_tensor, self._input_image_size + [self._num_channels]
      )
    return image_tensor
//  [python]
def __call__(self, observation, actions=None):
    """Returns the probability of input actions being feasible."""
    if actions is None:
      actions = tf.range(
          self._action_spec.minimum, self._action_spec.maximum + 1
      )
      actions = tf.reshape(actions, [1, -1])
      actions = tf.tile(actions, [self.batch_size, 1])
    feasibility_prob = tf.to_float(tf.greater(actions, self.threshold))
    return feasibility_prob
//  [python]
def __call__(self, observation, actions=None):
    """Returns the probability of input actions being feasible."""
    if actions is None:
      actions = tf.range(
          self._action_spec.minimum, self._action_spec.maximum + 1
      )
      actions = tf.reshape(actions, [1, -1])
      actions = tf.tile(actions, [self.batch_size, 1])
    feasibility_prob = tf.cast(tf.greater(actions, self.threshold), tf.float32)
    return feasibility_prob
//  [python]
def _create_test_trajectory(self, batch_size):
    num_actions = t.cast(batch_size / 2, dtype=t.int32)
    action_tensor = tf.concat(
        [
            tf.range(num_actions, dtype=tf.int32),
            tf.range(num_actions, dtype=tf.int32),
        ],
        axis=-1,
    )
    return trajectory.Trajectory(
        observation=tf.ones(batch_size),
        action=action_tensor,
        policy_info=(),
        reward=tf.range(batch_size, dtype=tf.float32),
        discount=tf.ones(batch_size),
        step_type=ts.StepType.FIRST,
        next_step_type=ts.StepType.LAST,
    )
//  [python]
def _create_test_trajectory(self, batch_size):
    num_actions = tf.cast(batch_size / 2, dtype=tf.int32)
    action_tensor = tf.concat(
        [
            tf.range(num_actions, dtype=tf.int32),
            tf.range(num_actions, dtype=tf.int32),
        ],
        axis=-1,
    )
    return trajectory.Trajectory(
        observation=tf.ones(batch_size),
        action=action_tensor,
        policy_info=(),
        reward=tf.range(batch_size, dtype=tf.float32),
        discount=tf.ones(batch_size),
        step_type=ts.StepType.FIRST,
        next_step_type=ts.StepType.LAST,
    )
//  [python]
def _filter_and_allreduce_gradients(grads_and_vars,
                                    allreduce_precision="float32",
                                    bytes_per_pack=0):
  """Filter None grads and then allreduce gradients in specified precision.

  This utils function is used when users intent to explicitly allreduce
  gradients and customize gradients operations before and after allreduce.
  The allreduced gradients are then passed to optimizer.apply_gradients(
  experimental_aggregate_gradients=False).

  Args:
      grads_and_vars: gradients and variables pairs.
      allreduce_precision: Whether to allreduce gradients in float32 or float16.
      bytes_per_pack: A non-negative integer. Breaks collective operations into
        packs of certain size. If it's zero, all gradients are in one pack.

  Returns:
      pairs of allreduced non-None gradients and variables.
  """
  filtered_grads_and_vars = _filter_grads(grads_and_vars)
  (grads, variables) = zip(*filtered_grads_and_vars)
  if allreduce_precision == "float16":
    grads = [t.cast(grad, "float16") for grad in grads]
  hints = tf.distribute.experimental.CommunicationOptions(
      bytes_per_pack=bytes_per_pack)
  allreduced_grads = tf.distribute.get_strategy(  # pylint: disable=protected-access
  ).extended._replica_ctx_all_reduce(tf.distribute.ReduceOp.SUM, grads, hints)
  if allreduce_precision == "float16":
    allreduced_grads = [t.cast(grad, "float32") for grad in allreduced_grads]
  return allreduced_grads, variables
//  [python]
def _filter_and_allreduce_gradients(grads_and_vars,
                                    allreduce_precision="float32",
                                    bytes_per_pack=0):
  """Filter None grads and then allreduce gradients in specified precision.

  This utils function is used when users intent to explicitly allreduce
  gradients and customize gradients operations before and after allreduce.
  The allreduced gradients are then passed to optimizer.apply_gradients(
  experimental_aggregate_gradients=False).

  Args:
      grads_and_vars: gradients and variables pairs.
      allreduce_precision: Whether to allreduce gradients in float32 or float16.
      bytes_per_pack: A non-negative integer. Breaks collective operations into
        packs of certain size. If it's zero, all gradients are in one pack.

  Returns:
      pairs of allreduced non-None gradients and variables.
  """
  filtered_grads_and_vars = _filter_grads(grads_and_vars)
  (grads, variables) = zip(*filtered_grads_and_vars)
  if allreduce_precision == "float16":
    grads = [tf.cast(grad, "float16") for grad in grads]
  hints = tf.distribute.experimental.CommunicationOptions(
      bytes_per_pack=bytes_per_pack)
  allreduced_grads = tf.distribute.get_strategy(  # pylint: disable=protected-access
  ).extended._replica_ctx_all_reduce(tf.distribute.ReduceOp.SUM, grads, hints)
  if allreduce_precision == "float16":
    allreduced_grads = [tf.cast(grad, "float32") for grad in allreduced_grads]
  return allreduced_grads, variables
//  [python]
def gen_cast_test(name, input_shape, dtype):
    # Create model.
    inp = layers.Input(name="input", batch_size=input_shape[0], shape=input_shape[1:])
    out = tf.cast(inp, dtype=dtype)
    model = Model(inputs=[inp], outputs=[out])
    # Create data.
    np.random.seed(0)
    inp_tensor = np.random.rand(*input_shape).astype(np.float32)
    out_tensor = model.predict(inp_tensor)
    # Save model.
    save_model(model, name)
    # Save data.
    save_tensor(inp_tensor, name + ".inp0")
    save_tensor(out_tensor, name + ".out0")
    # Clear session.
    keras_backend.clear_session()