jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_gcs_to_lcs(self):
        """Test 3GPP channel coefficient calculation: GCS to LCS"""
        batch_size = TestChannelCoefficientsGenerator.BATCH_SIZE
        orientation = tf.random.normal(shape=[batch_size,3]).numpy()
        theta = tf.random.normal(shape=[batch_size]).numpy()
        phi = tf.random.normal(shape=[batch_size]).numpy()

        theta_prime_ref, phi_prime_ref = self.gcs_to_lcs_ref(orientation, theta,
                                                            phi)
        theta_prime, phi_prime = self.ccg._gcs_to_lcs(
            t.cast(orientation, tf.float64),
            t.cast(theta, tf.float64),
            t.cast(phi, tf.float64))
        theta_prime = theta_prime.numpy()
        phi_prime = phi_prime.numpy()

        err_tol = TestChannelCoefficientsGenerator.MAX_ERR
        max_err = self.max_rel_err(theta_prime_ref, theta_prime)
        self.assertLessEqual(max_err, err_tol)
        max_err = self.max_rel_err(phi_prime_ref, phi_prime)
        self.assertLessEqual(max_err, err_tol)
//  [python]
def test_gcs_to_lcs(self):
        """Test 3GPP channel coefficient calculation: GCS to LCS"""
        batch_size = TestChannelCoefficientsGenerator.BATCH_SIZE
        orientation = tf.random.normal(shape=[batch_size,3]).numpy()
        theta = tf.random.normal(shape=[batch_size]).numpy()
        phi = tf.random.normal(shape=[batch_size]).numpy()

        theta_prime_ref, phi_prime_ref = self.gcs_to_lcs_ref(orientation, theta,
                                                            phi)
        theta_prime, phi_prime = self.ccg._gcs_to_lcs(
            tf.cast(orientation, tf.float64),
            tf.cast(theta, tf.float64),
            tf.cast(phi, tf.float64))
        theta_prime = theta_prime.numpy()
        phi_prime = phi_prime.numpy()

        err_tol = TestChannelCoefficientsGenerator.MAX_ERR
        max_err = self.max_rel_err(theta_prime_ref, theta_prime)
        self.assertLessEqual(max_err, err_tol)
        max_err = self.max_rel_err(phi_prime_ref, phi_prime)
        self.assertLessEqual(max_err, err_tol)
//  [python]
def test_l2g_response(self):
        """Test 3GPP channel coefficient calculation: L2G antenna response"""
        batch_size = TestChannelCoefficientsGenerator.BATCH_SIZE
        orientation = tf.random.normal(shape=[batch_size,3]).numpy()
        theta = tf.random.normal(shape=[batch_size]).numpy()
        phi = tf.random.normal(shape=[batch_size]).numpy()
        F_prime = tf.random.normal(shape=[batch_size,2]).numpy()

        F_ref = self.l2g_response_ref(F_prime, orientation, theta, phi)
        F = self.ccg._l2g_response( t.cast(F_prime, tf.float64),
                                    t.cast(orientation,tf.float64),
                                    t.cast(theta, tf.float64),
                                    t.cast(phi, tf.float64)).numpy()

        max_err = self.max_rel_err(F_ref, F)
        err_tol = TestChannelCoefficientsGenerator.MAX_ERR
        self.assertLessEqual(max_err, err_tol)
//  [python]
def test_l2g_response(self):
        """Test 3GPP channel coefficient calculation: L2G antenna response"""
        batch_size = TestChannelCoefficientsGenerator.BATCH_SIZE
        orientation = tf.random.normal(shape=[batch_size,3]).numpy()
        theta = tf.random.normal(shape=[batch_size]).numpy()
        phi = tf.random.normal(shape=[batch_size]).numpy()
        F_prime = tf.random.normal(shape=[batch_size,2]).numpy()

        F_ref = self.l2g_response_ref(F_prime, orientation, theta, phi)
        F = self.ccg._l2g_response( tf.cast(F_prime, tf.float64),
                                    tf.cast(orientation,tf.float64),
                                    tf.cast(theta, tf.float64),
                                    tf.cast(phi, tf.float64)).numpy()

        max_err = self.max_rel_err(F_ref, F)
        err_tol = TestChannelCoefficientsGenerator.MAX_ERR
        self.assertLessEqual(max_err, err_tol)
//  [python]
def _step_11_nlos(self, phi, topology, rays, t):
        # pylint: disable=line-too-long
        r"""
        Compute the full NLOS channel matrix (7.5-28)

        Input
        -----
        phi: [batch size, num TXs, num RXs, num clusters, num rays, 4], tf.float
            Random initial phases [radian]

        topology : Topology
            Topology of the network

        rays : Rays
            Rays

        t : [num time samples], tf.float
            Time samples

        Output
        ------
        h_full : [batch size, num_tx, num rx, num clusters, num rays, num rx antennas, num tx antennas, num time steps], tf.complex
            NLoS channel matrix
        """

        h_phase = self._step_11_phase_matrix(phi, rays)
        h_field = self._step_11_field_matrix(topology, rays.aoa, rays.aod,
                                                    rays.zoa, rays.zod, h_phase)
        h_array = self._step_11_array_offsets(topology, rays.aoa, rays.aod,
                                                            rays.zoa, rays.zod)
        h_doppler = self._step_11_doppler_matrix(topology, rays.aoa, rays.zoa,
                                                                            t)
        h_full = tf.expand_dims(h_field*h_array, -1) * tf.expand_dims(
            tf.expand_dims(h_doppler, -2), -2)

        power_scaling = tf.complex(tf.sqrt(rays.powers/
            tf.to_float(tf.shape(h_full)[4]), self._dtype.real_dtype),
                            tf.constant(0., self._dtype.real_dtype))
        shape = tf.concat([tf.shape(power_scaling), tf.ones(
            [tf.rank(h_full)-tf.rank(power_scaling)], tf.int32)], 0)
        h_full *= tf.reshape(power_scaling, shape)

        return h_full
//  [python]
def _step_11_nlos(self, phi, topology, rays, t):
        # pylint: disable=line-too-long
        r"""
        Compute the full NLOS channel matrix (7.5-28)

        Input
        -----
        phi: [batch size, num TXs, num RXs, num clusters, num rays, 4], tf.float
            Random initial phases [radian]

        topology : Topology
            Topology of the network

        rays : Rays
            Rays

        t : [num time samples], tf.float
            Time samples

        Output
        ------
        h_full : [batch size, num_tx, num rx, num clusters, num rays, num rx antennas, num tx antennas, num time steps], tf.complex
            NLoS channel matrix
        """

        h_phase = self._step_11_phase_matrix(phi, rays)
        h_field = self._step_11_field_matrix(topology, rays.aoa, rays.aod,
                                                    rays.zoa, rays.zod, h_phase)
        h_array = self._step_11_array_offsets(topology, rays.aoa, rays.aod,
                                                            rays.zoa, rays.zod)
        h_doppler = self._step_11_doppler_matrix(topology, rays.aoa, rays.zoa,
                                                                            t)
        h_full = tf.expand_dims(h_field*h_array, -1) * tf.expand_dims(
            tf.expand_dims(h_doppler, -2), -2)

        power_scaling = tf.complex(tf.sqrt(rays.powers/
            tf.cast(tf.shape(h_full)[4], self._dtype.real_dtype)),
                            tf.constant(0., self._dtype.real_dtype))
        shape = tf.concat([tf.shape(power_scaling), tf.ones(
            [tf.rank(h_full)-tf.rank(power_scaling)], tf.int32)], 0)
        h_full *= tf.reshape(power_scaling, shape)

        return h_full
//  [python]
def _build_target_q_op(self):
    """Build an op used as a target for the Q-value.

    Returns:
      target_q_op: An op calculating the Q-value.
    """
    # Get the maximum Q-value across the actions dimension.
    replay_next_qt_max = tf.reduce_max(
        self._replay_next_target_net_outputs.q_values, 1
    )
    # Calculate the Bellman target value.
    #   Q_t = R_t + \gamma^N * Q'_t+1
    # where,
    #   Q'_t+1 = \argmax_a Q(S_t+1, a)
    #          (or) 0 if S_t is a terminal state,
    # and
    #   N is the update horizon (by default, N=1).
    return self._replay.rewards + self.cumulative_gamma * replay_next_qt_max * (
        1.0 - tf.to_float(self._replay.terminals)
    )
//  [python]
def _build_target_q_op(self):
    """Build an op used as a target for the Q-value.

    Returns:
      target_q_op: An op calculating the Q-value.
    """
    # Get the maximum Q-value across the actions dimension.
    replay_next_qt_max = tf.reduce_max(
        self._replay_next_target_net_outputs.q_values, 1
    )
    # Calculate the Bellman target value.
    #   Q_t = R_t + \gamma^N * Q'_t+1
    # where,
    #   Q'_t+1 = \argmax_a Q(S_t+1, a)
    #          (or) 0 if S_t is a terminal state,
    # and
    #   N is the update horizon (by default, N=1).
    return self._replay.rewards + self.cumulative_gamma * replay_next_qt_max * (
        1.0 - tf.cast(self._replay.terminals, tf.float32)
    )
//  [python]
def _get_learning_rate(self, step):
    """Compute learning rate at given step."""
    step = tf.cast(step, dtype=tf.float32)
    warmup_steps = tf.cast(self.warmup_steps, dtype=tf.float32)
    with tf.name_scope('PiecewiseConstantDecayWithWarmup'):

      def warmup_lr(step):
        return self.rescaled_lr * (step / warmup_steps)

      def piecewise_lr(step):
        return tf.compat.v1.train.piecewise_constant(step, self.step_boundaries,
                                                     self.lr_values)

      return tf.cond(step < warmup_steps, lambda: warmup_lr(step),
                     lambda: piecewise_lr(step))
//  [python]
def visualize_images_with_bounding_boxes(images, box_outputs, step,
                                         summary_writer):
  """Records subset of evaluation images with bounding boxes."""
  image_shape = tf.shape(images[0])
  image_height = tf.to_float(image_shape[0])
  image_width = tf.to_float(image_shape[1])
  normalized_boxes = normalize_boxes(box_outputs, [image_height, image_width])

  bounding_box_color = tf.constant([[1.0, 1.0, 0.0, 1.0]])
  image_summary = tf.image.draw_bounding_boxes(images, normalized_boxes,
                                               bounding_box_color)
  with summary_writer.as_default():
    tf.summary.image('bounding_box_summary', image_summary, step=step)
    summary_writer.flush()
//  [python]
def visualize_images_with_bounding_boxes(images, box_outputs, step,
                                         summary_writer):
  """Records subset of evaluation images with bounding boxes."""
  image_shape = tf.shape(images[0])
  image_height = tf.cast(image_shape[0], tf.float32)
  image_width = tf.cast(image_shape[1], tf.float32)
  normalized_boxes = normalize_boxes(box_outputs, [image_height, image_width])

  bounding_box_color = tf.constant([[1.0, 1.0, 0.0, 1.0]])
  image_summary = tf.image.draw_bounding_boxes(images, normalized_boxes,
                                               bounding_box_color)
  with summary_writer.as_default():
    tf.summary.image('bounding_box_summary', image_summary, step=step)
    summary_writer.flush()
//  [python]
def check_labels_correct(self,
                           boxes,
                           classes,
                           output_size,
                           input_size):
    max_num_instances = 128
    num_detections = len(boxes)
    boxes = tf.constant(boxes, dtype=tf.float32)
    classes = tf.constant(classes, dtype=tf.float32)

    boxes = preprocess_ops.clip_or_pad_to_fixed_size(
        boxes, max_num_instances, 0)
    classes = preprocess_ops.clip_or_pad_to_fixed_size(
        classes, max_num_instances, 0)

    # pylint: disable=g-long-lambda
    labels = target_assigner.assign_centernet_targets(
        labels={
            'boxes': boxes,
            'classes': classes,
            'groundtruths': {
                'num_detections': num_detections,
            }
        },
        output_size=output_size,
        input_size=input_size)

    ct_heatmaps = labels['ct_heatmaps']
    ct_offset = labels['ct_offset']
    size = labels['size']
    box_mask = labels['box_mask']
    box_indices = labels['box_indices']

    boxes = tf.cast(boxes, tf.float32)
    classes = tf.cast(classes, tf.float32)
    height_ratio = output_size[0] / input_size[0]
    width_ratio = output_size[1] / input_size[1]

    # Shape checks
    self.assertEqual(ct_heatmaps.shape, (output_size[0], output_size[1], 90))

    self.assertEqual(ct_offset.shape, (max_num_instances, 2))

    self.assertEqual(size.shape, (max_num_instances, 2))
    self.assertEqual(box_mask.shape, (max_num_instances,))
    self.assertEqual(box_indices.shape, (max_num_instances, 2))

    self.assertAllInRange(ct_heatmaps, 0, 1)

    for i in range(len(boxes)):
      # Check sizes
      self.assertAllEqual(size[i],
                          [(boxes[i][2] - boxes[i][0]) * height_ratio,
                           (boxes[i][3] - boxes[i][1]) * width_ratio,
                           ])

      # Check box indices
      y = tf.math.floor((boxes[i][0] + boxes[i][2]) / 2 * height_ratio)
      x = tf.math.floor((boxes[i][1] + boxes[i][3]) / 2 * width_ratio)
      self.assertAllEqual(box_indices[i], [y, x])

      # check offsets
      true_y = (boxes[i][0] + boxes[i][2]) / 2 * height_ratio
      true_x = (boxes[i][1] + boxes[i][3]) / 2 * width_ratio
      self.assertAllEqual(ct_offset[i], [true_y - y, true_x - x])

    for i in range(len(boxes), max_num_instances):
      # Make sure rest are zero
      self.assertAllEqual(size[i], [0, 0])
      self.assertAllEqual(box_indices[i], [0, 0])
      self.assertAllEqual(ct_offset[i], [0, 0])

    # Check mask indices
    self.assertAllEqual(tf.cast(box_mask[3:], tf.int32),
                        tf.repeat(0, repeats=max_num_instances - 3))
    self.assertAllEqual(tf.cast(box_mask[:3], tf.int32),
                        tf.repeat(1, repeats=3))
//  [python]
def denormalize_boxes(boxes, image_shape):
  """Converts boxes normalized by [height, width] to pixel coordinates.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    denormalized_boxes: a tensor whose shape is the same as `boxes` representing
      the denormalized boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  with tf.name_scope('denormalize_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
    else:
      image_shape = tf.to_float(image_shape)
      height, width = tf.split(image_shape, 2, axis=-1)

    ymin, xmin, ymax, xmax = tf.split(boxes, 4, axis=-1)
    ymin = ymin * height
    xmin = xmin * width
    ymax = ymax * height
    xmax = xmax * width

    denormalized_boxes = tf.concat([ymin, xmin, ymax, xmax], axis=-1)
    return denormalized_boxes
//  [python]
def denormalize_boxes(boxes, image_shape):
  """Converts boxes normalized by [height, width] to pixel coordinates.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    denormalized_boxes: a tensor whose shape is the same as `boxes` representing
      the denormalized boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  with tf.name_scope('denormalize_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
    else:
      image_shape = tf.cast(image_shape, dtype=boxes.dtype)
      height, width = tf.split(image_shape, 2, axis=-1)

    ymin, xmin, ymax, xmax = tf.split(boxes, 4, axis=-1)
    ymin = ymin * height
    xmin = xmin * width
    ymax = ymax * height
    xmax = xmax * width

    denormalized_boxes = tf.concat([ymin, xmin, ymax, xmax], axis=-1)
    return denormalized_boxes
//  [python]
def normalize_boxes(boxes, image_shape):
  """Converts boxes to the normalized coordinates.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    normalized_boxes: a tensor whose shape is the same as `boxes` representing
      the normalized boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError('boxes.shape[-1] is {:d}, but must be 4.'.format(
        boxes.shape[-1]))

  with tf.name_scope('normalize_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
    else:
      image_shape = tf.to_float(image_shape)
      height = image_shape[..., 0:1]
      width = image_shape[..., 1:2]

    ymin = boxes[..., 0:1] / height
    xmin = boxes[..., 1:2] / width
    ymax = boxes[..., 2:3] / height
    xmax = boxes[..., 3:4] / width

    normalized_boxes = tf.concat([ymin, xmin, ymax, xmax], axis=-1)
    return normalized_boxes
//  [python]
def normalize_boxes(boxes, image_shape):
  """Converts boxes to the normalized coordinates.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    normalized_boxes: a tensor whose shape is the same as `boxes` representing
      the normalized boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError('boxes.shape[-1] is {:d}, but must be 4.'.format(
        boxes.shape[-1]))

  with tf.name_scope('normalize_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
    else:
      image_shape = tf.cast(image_shape, dtype=boxes.dtype)
      height = image_shape[..., 0:1]
      width = image_shape[..., 1:2]

    ymin = boxes[..., 0:1] / height
    xmin = boxes[..., 1:2] / width
    ymax = boxes[..., 2:3] / height
    xmax = boxes[..., 3:4] / width

    normalized_boxes = tf.concat([ymin, xmin, ymax, xmax], axis=-1)
    return normalized_boxes
//  [python]
def get_synth_input_fn(height,
                       width,
                       num_channels,
                       num_classes,
                       dtype=tf.float32,
                       drop_remainder=True):
  """Returns an input function that returns a dataset with random data.

  This input_fn returns a data set that iterates over a set of random data and
  bypasses all preprocessing, e.g. jpeg decode and copy. The host to device
  copy is still included. This used to find the upper throughput bound when
  tuning the full input pipeline.

  Args:
    height: Integer height that will be used to create a fake image tensor.
    width: Integer width that will be used to create a fake image tensor.
    num_channels: Integer depth that will be used to create a fake image tensor.
    num_classes: Number of classes that should be represented in the fake labels
      tensor
    dtype: Data type for features/images.
    drop_remainder: A boolean indicates whether to drop the remainder of the
      batches. If True, the batch dimension will be static.

  Returns:
    An input_fn that can be used in place of a real one to return a dataset
    that can be used for iteration.
  """

  # pylint: disable=unused-argument
  def input_fn(is_training, data_dir, batch_size, *args, **kwargs):
    """Returns dataset filled with random data."""
    inputs, labels = get_synth_data(
        height=height,
        width=width,
        num_channels=num_channels,
        num_classes=num_classes,
        dtype=dtype)
    # Cast to float32 for Keras model.
    labels = tf.to_float(labels)
    data = tf.data.Dataset.from_tensors((inputs, labels)).repeat()

    # `drop_remainder` will make dataset produce outputs with known shapes.
    data = data.batch(batch_size, drop_remainder=drop_remainder)
    data = data.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
    return data

  return input_fn
//  [python]
def get_synth_input_fn(height,
                       width,
                       num_channels,
                       num_classes,
                       dtype=tf.float32,
                       drop_remainder=True):
  """Returns an input function that returns a dataset with random data.

  This input_fn returns a data set that iterates over a set of random data and
  bypasses all preprocessing, e.g. jpeg decode and copy. The host to device
  copy is still included. This used to find the upper throughput bound when
  tuning the full input pipeline.

  Args:
    height: Integer height that will be used to create a fake image tensor.
    width: Integer width that will be used to create a fake image tensor.
    num_channels: Integer depth that will be used to create a fake image tensor.
    num_classes: Number of classes that should be represented in the fake labels
      tensor
    dtype: Data type for features/images.
    drop_remainder: A boolean indicates whether to drop the remainder of the
      batches. If True, the batch dimension will be static.

  Returns:
    An input_fn that can be used in place of a real one to return a dataset
    that can be used for iteration.
  """

  # pylint: disable=unused-argument
  def input_fn(is_training, data_dir, batch_size, *args, **kwargs):
    """Returns dataset filled with random data."""
    inputs, labels = get_synth_data(
        height=height,
        width=width,
        num_channels=num_channels,
        num_classes=num_classes,
        dtype=dtype)
    # Cast to float32 for Keras model.
    labels = tf.cast(labels, dtype=tf.float32)
    data = tf.data.Dataset.from_tensors((inputs, labels)).repeat()

    # `drop_remainder` will make dataset produce outputs with known shapes.
    data = data.batch(batch_size, drop_remainder=drop_remainder)
    data = data.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
    return data

  return input_fn
//  [python]
def clip_boxes(boxes, image_shape):
  """Clips boxes to image boundaries.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    clipped_boxes: a tensor whose shape is the same as `boxes` representing the
      clipped boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError('boxes.shape[-1] is {:d}, but must be 4.'.format(
        boxes.shape[-1]))

  with tf.name_scope('clip_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
      max_length = [height - 1.0, width - 1.0, height - 1.0, width - 1.0]
    else:
      image_shape = tf.to_float(image_shape)
      height, width = tf.unstack(image_shape, axis=-1)
      max_length = tf.stack(
          [height - 1.0, width - 1.0, height - 1.0, width - 1.0], axis=-1)

    clipped_boxes = tf.math.maximum(tf.math.minimum(boxes, max_length), 0.0)
    return clipped_boxes
//  [python]
def clip_boxes(boxes, image_shape):
  """Clips boxes to image boundaries.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    image_shape: a list of two integers, a two-element vector or a tensor such
      that all but the last dimensions are `broadcastable` to `boxes`. The last
      dimension is 2, which represents [height, width].

  Returns:
    clipped_boxes: a tensor whose shape is the same as `boxes` representing the
      clipped boxes.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError('boxes.shape[-1] is {:d}, but must be 4.'.format(
        boxes.shape[-1]))

  with tf.name_scope('clip_boxes'):
    if isinstance(image_shape, list) or isinstance(image_shape, tuple):
      height, width = image_shape
      max_length = [height - 1.0, width - 1.0, height - 1.0, width - 1.0]
    else:
      image_shape = tf.cast(image_shape, dtype=boxes.dtype)
      height, width = tf.unstack(image_shape, axis=-1)
      max_length = tf.stack(
          [height - 1.0, width - 1.0, height - 1.0, width - 1.0], axis=-1)

    clipped_boxes = tf.math.maximum(tf.math.minimum(boxes, max_length), 0.0)
    return clipped_boxes
//  [python]
def filter_boxes_by_scores(boxes, scores, min_score_threshold):
  """Filter and remove boxes whose scores are smaller than the threshold.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    scores: a tensor whose shape is the same as tf.shape(boxes)[:-1]
      representing the original scores of the boxes.
    min_score_threshold: a float representing the minimal box score threshold.
      Boxes whose score are smaller than it will be filtered out.

  Returns:
    filtered_boxes: a tensor whose shape is the same as `boxes` but with
      the position of the filtered boxes are filled with -1.
    filtered_scores: a tensor whose shape is the same as 'scores' but with
      the
  """
  if boxes.shape[-1] != 4:
    raise ValueError('boxes.shape[1] is {:d}, but must be 4.'.format(
        boxes.shape[-1]))

  with tf.name_scope('filter_boxes_by_scores'):
    filtered_mask = tf.math.greater(scores, min_score_threshold)
    filtered_scores = tf.where(filtered_mask, scores, -tf.ones_like(scores))
    filtered_boxes = tf.to_float(
        tf.expand_dims(filtered_mask, axis=-1)) * boxes

    return filtered_boxes, filtered_scores
//  [python]
def filter_boxes_by_scores(boxes, scores, min_score_threshold):
  """Filter and remove boxes whose scores are smaller than the threshold.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    scores: a tensor whose shape is the same as tf.shape(boxes)[:-1]
      representing the original scores of the boxes.
    min_score_threshold: a float representing the minimal box score threshold.
      Boxes whose score are smaller than it will be filtered out.

  Returns:
    filtered_boxes: a tensor whose shape is the same as `boxes` but with
      the position of the filtered boxes are filled with -1.
    filtered_scores: a tensor whose shape is the same as 'scores' but with
      the
  """
  if boxes.shape[-1] != 4:
    raise ValueError('boxes.shape[1] is {:d}, but must be 4.'.format(
        boxes.shape[-1]))

  with tf.name_scope('filter_boxes_by_scores'):
    filtered_mask = tf.math.greater(scores, min_score_threshold)
    filtered_scores = tf.where(filtered_mask, scores, -tf.ones_like(scores))
    filtered_boxes = tf.cast(
        tf.expand_dims(filtered_mask, axis=-1), dtype=boxes.dtype) * boxes

    return filtered_boxes, filtered_scores
//  [python]
def encode_boxes(boxes, anchors, weights=None):
  """Encode boxes to targets.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    anchors: a tensor whose shape is the same as, or `broadcastable` to `boxes`,
      representing the coordinates of anchors in ymin, xmin, ymax, xmax order.
    weights: None or a list of four float numbers used to scale coordinates.

  Returns:
    encoded_boxes: a tensor whose shape is the same as `boxes` representing the
      encoded box targets.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError('boxes.shape[-1] is {:d}, but must be 4.'.format(
        boxes.shape[-1]))

  with tf.name_scope('encode_boxes'):
    boxes = tf.to_float(boxes, dtype=anchors.dtype)
    ymin = boxes[..., 0:1]
    xmin = boxes[..., 1:2]
    ymax = boxes[..., 2:3]
    xmax = boxes[..., 3:4]
    box_h = ymax - ymin + 1.0
    box_w = xmax - xmin + 1.0
    box_yc = ymin + 0.5 * box_h
    box_xc = xmin + 0.5 * box_w

    anchor_ymin = anchors[..., 0:1]
    anchor_xmin = anchors[..., 1:2]
    anchor_ymax = anchors[..., 2:3]
    anchor_xmax = anchors[..., 3:4]
    anchor_h = anchor_ymax - anchor_ymin + 1.0
    anchor_w = anchor_xmax - anchor_xmin + 1.0
    anchor_yc = anchor_ymin + 0.5 * anchor_h
    anchor_xc = anchor_xmin + 0.5 * anchor_w

    encoded_dy = (box_yc - anchor_yc) / anchor_h
    encoded_dx = (box_xc - anchor_xc) / anchor_w
    encoded_dh = tf.math.log(box_h / anchor_h)
    encoded_dw = tf.math.log(box_w / anchor_w)
    if weights:
      encoded_dy *= weights[0]
      encoded_dx *= weights[1]
      encoded_dh *= weights[2]
      encoded_dw *= weights[3]

    encoded_boxes = tf.concat([encoded_dy, encoded_dx, encoded_dh, encoded_dw],
                              axis=-1)
    return encoded_boxes
//  [python]
def encode_boxes(boxes, anchors, weights=None):
  """Encode boxes to targets.

  Args:
    boxes: a tensor whose last dimension is 4 representing the coordinates of
      boxes in ymin, xmin, ymax, xmax order.
    anchors: a tensor whose shape is the same as, or `broadcastable` to `boxes`,
      representing the coordinates of anchors in ymin, xmin, ymax, xmax order.
    weights: None or a list of four float numbers used to scale coordinates.

  Returns:
    encoded_boxes: a tensor whose shape is the same as `boxes` representing the
      encoded box targets.

  Raises:
    ValueError: If the last dimension of boxes is not 4.
  """
  if boxes.shape[-1] != 4:
    raise ValueError('boxes.shape[-1] is {:d}, but must be 4.'.format(
        boxes.shape[-1]))

  with tf.name_scope('encode_boxes'):
    boxes = tf.cast(boxes, dtype=anchors.dtype)
    ymin = boxes[..., 0:1]
    xmin = boxes[..., 1:2]
    ymax = boxes[..., 2:3]
    xmax = boxes[..., 3:4]
    box_h = ymax - ymin + 1.0
    box_w = xmax - xmin + 1.0
    box_yc = ymin + 0.5 * box_h
    box_xc = xmin + 0.5 * box_w

    anchor_ymin = anchors[..., 0:1]
    anchor_xmin = anchors[..., 1:2]
    anchor_ymax = anchors[..., 2:3]
    anchor_xmax = anchors[..., 3:4]
    anchor_h = anchor_ymax - anchor_ymin + 1.0
    anchor_w = anchor_xmax - anchor_xmin + 1.0
    anchor_yc = anchor_ymin + 0.5 * anchor_h
    anchor_xc = anchor_xmin + 0.5 * anchor_w

    encoded_dy = (box_yc - anchor_yc) / anchor_h
    encoded_dx = (box_xc - anchor_xc) / anchor_w
    encoded_dh = tf.math.log(box_h / anchor_h)
    encoded_dw = tf.math.log(box_w / anchor_w)
    if weights:
      encoded_dy *= weights[0]
      encoded_dx *= weights[1]
      encoded_dh *= weights[2]
      encoded_dw *= weights[3]

    encoded_boxes = tf.concat([encoded_dy, encoded_dx, encoded_dh, encoded_dw],
                              axis=-1)
    return encoded_boxes
//  [python]
def top_k_with_unique(inputs, k):
  """Finds the values and indices of the k largests entries.
  Instead of doing sort like tf.nn.top_k, this function finds the max value
  k times. The running time is proportional to k, which is be faster when k
  is small. The current implementation supports only inputs of rank 2.
  In addition, iota is used to replace the lower bits of each element, this
  makes the selection more stable when there are equal elements. The
  overhead is that output values are approximated.
  Args:
    inputs: A tensor with rank of 2. [batch_size, original_size].
    k: An integer, number of top elements to select.
  Returns:
    top_values: A tensor, the k largest elements in sorted order.
      [batch_size, k].
    indices: A tensor, indices of the top_values. [batch_size, k].
  """
  unique_inputs = _create_make_unique(t.cast(inputs, t.float32))
  top_values, indices = _create_topk_unique(unique_inputs, k)
  top_values = t.cast(top_values, inputs.dtype)
  return top_values, indices
//  [python]
def top_k_with_unique(inputs, k):
  """Finds the values and indices of the k largests entries.
  Instead of doing sort like tf.nn.top_k, this function finds the max value
  k times. The running time is proportional to k, which is be faster when k
  is small. The current implementation supports only inputs of rank 2.
  In addition, iota is used to replace the lower bits of each element, this
  makes the selection more stable when there are equal elements. The
  overhead is that output values are approximated.
  Args:
    inputs: A tensor with rank of 2. [batch_size, original_size].
    k: An integer, number of top elements to select.
  Returns:
    top_values: A tensor, the k largest elements in sorted order.
      [batch_size, k].
    indices: A tensor, indices of the top_values. [batch_size, k].
  """
  unique_inputs = _create_make_unique(tf.cast(inputs, tf.float32))
  top_values, indices = _create_topk_unique(unique_inputs, k)
  top_values = tf.cast(top_values, inputs.dtype)
  return top_values, indices
//  [python]
def project_svd(basis, vector, singular_threshold):
  """Project `vector` into the colspace of `basis' using SVD.

  Args:
    basis: a column space into which we will project, shape [B, N, M]
    vector: a vector to be projected, shape [B, N, 1]
    singular_threshold: singular values below this threshold are treated as 0

  Returns:
    projected_vector: the same shape as vector
    weights: [B, M]: weights so that basis*vector=projected_vector. If basis
                     is not full-rank, this is non-unique and no guarantees are
                     made about which set is returned.
  """
  s, u, v = tf.linalg.svd(basis)
  # The columns of U form an orthonormal basis and the matrix 'subspace'
  # indicates a subset of those which form an orthonormal basis for the column
  # space of M.
  subspace = tf.abs(s) >= singular_threshold
  subspace = tf.to_float(subspace)

  projected_vector = tf.matmul(
      u,
      tf.matmul(
          tf.linalg.diag(subspace),
          tf.matmul(u, vector, transpose_a=True)))

  # We can find weights that give this projected vector by solving the normal
  # equations (A^TA)^(-1)A^Tx.

  # 1/s where s > singular_threshold, 0 elsewhere
  s_inverse = tf.math.divide_no_nan(subspace, s)
  weights = tf.matmul(
      v,
      tf.matmul(
          tf.linalg.diag(s_inverse),
          tf.matmul(u, vector, transpose_a=True)))
  return projected_vector, tf.squeeze(weights, -1)
//  [python]
def project_svd(basis, vector, singular_threshold):
  """Project `vector` into the colspace of `basis' using SVD.

  Args:
    basis: a column space into which we will project, shape [B, N, M]
    vector: a vector to be projected, shape [B, N, 1]
    singular_threshold: singular values below this threshold are treated as 0

  Returns:
    projected_vector: the same shape as vector
    weights: [B, M]: weights so that basis*vector=projected_vector. If basis
                     is not full-rank, this is non-unique and no guarantees are
                     made about which set is returned.
  """
  s, u, v = tf.linalg.svd(basis)
  # The columns of U form an orthonormal basis and the matrix 'subspace'
  # indicates a subset of those which form an orthonormal basis for the column
  # space of M.
  subspace = tf.abs(s) >= singular_threshold
  subspace = tf.cast(subspace, tf.float32)

  projected_vector = tf.matmul(
      u,
      tf.matmul(
          tf.linalg.diag(subspace),
          tf.matmul(u, vector, transpose_a=True)))

  # We can find weights that give this projected vector by solving the normal
  # equations (A^TA)^(-1)A^Tx.

  # 1/s where s > singular_threshold, 0 elsewhere
  s_inverse = tf.math.divide_no_nan(subspace, s)
  weights = tf.matmul(
      v,
      tf.matmul(
          tf.linalg.diag(s_inverse),
          tf.matmul(u, vector, transpose_a=True)))
  return projected_vector, tf.squeeze(weights, -1)
//  [python]
def decode_boxes(encoded_boxes, anchors, weights=None):
  """Decode boxes.

  Args:
    encoded_boxes: a tensor whose last dimension is 4 representing the
      coordinates of encoded boxes in ymin, xmin, ymax, xmax order.
    anchors: a tensor whose shape is the same as, or `broadcastable` to `boxes`,
      representing the coordinates of anchors in ymin, xmin, ymax, xmax order.
    weights: None or a list of four float numbers used to scale coordinates.

  Returns:
    encoded_boxes: a tensor whose shape is the same as `boxes` representing the
      decoded box targets.
  """
  if encoded_boxes.shape[-1] != 4:
    raise ValueError('encoded_boxes.shape[-1] is {:d}, but must be 4.'.format(
        encoded_boxes.shape[-1]))

  with tf.name_scope('decode_boxes'):
    encoded_boxes = tf.to_float(encoded_boxes, dtype=anchors.dtype)
    dy = encoded_boxes[..., 0:1]
    dx = encoded_boxes[..., 1:2]
    dh = encoded_boxes[..., 2:3]
    dw = encoded_boxes[..., 3:4]
    if weights:
      dy /= weights[0]
      dx /= weights[1]
      dh /= weights[2]
      dw /= weights[3]
    dh = tf.math.minimum(dh, BBOX_XFORM_CLIP)
    dw = tf.math.minimum(dw, BBOX_XFORM_CLIP)

    anchor_ymin = anchors[..., 0:1]
    anchor_xmin = anchors[..., 1:2]
    anchor_ymax = anchors[..., 2:3]
    anchor_xmax = anchors[..., 3:4]
    anchor_h = anchor_ymax - anchor_ymin + 1.0
    anchor_w = anchor_xmax - anchor_xmin + 1.0
    anchor_yc = anchor_ymin + 0.5 * anchor_h
    anchor_xc = anchor_xmin + 0.5 * anchor_w

    decoded_boxes_yc = dy * anchor_h + anchor_yc
    decoded_boxes_xc = dx * anchor_w + anchor_xc
    decoded_boxes_h = tf.math.exp(dh) * anchor_h
    decoded_boxes_w = tf.math.exp(dw) * anchor_w

    decoded_boxes_ymin = decoded_boxes_yc - 0.5 * decoded_boxes_h
    decoded_boxes_xmin = decoded_boxes_xc - 0.5 * decoded_boxes_w
    decoded_boxes_ymax = decoded_boxes_ymin + decoded_boxes_h - 1.0
    decoded_boxes_xmax = decoded_boxes_xmin + decoded_boxes_w - 1.0

    decoded_boxes = tf.concat([
        decoded_boxes_ymin, decoded_boxes_xmin, decoded_boxes_ymax,
        decoded_boxes_xmax
    ],
                              axis=-1)
    return decoded_boxes
//  [python]
def decode_boxes(encoded_boxes, anchors, weights=None):
  """Decode boxes.

  Args:
    encoded_boxes: a tensor whose last dimension is 4 representing the
      coordinates of encoded boxes in ymin, xmin, ymax, xmax order.
    anchors: a tensor whose shape is the same as, or `broadcastable` to `boxes`,
      representing the coordinates of anchors in ymin, xmin, ymax, xmax order.
    weights: None or a list of four float numbers used to scale coordinates.

  Returns:
    encoded_boxes: a tensor whose shape is the same as `boxes` representing the
      decoded box targets.
  """
  if encoded_boxes.shape[-1] != 4:
    raise ValueError('encoded_boxes.shape[-1] is {:d}, but must be 4.'.format(
        encoded_boxes.shape[-1]))

  with tf.name_scope('decode_boxes'):
    encoded_boxes = tf.cast(encoded_boxes, dtype=anchors.dtype)
    dy = encoded_boxes[..., 0:1]
    dx = encoded_boxes[..., 1:2]
    dh = encoded_boxes[..., 2:3]
    dw = encoded_boxes[..., 3:4]
    if weights:
      dy /= weights[0]
      dx /= weights[1]
      dh /= weights[2]
      dw /= weights[3]
    dh = tf.math.minimum(dh, BBOX_XFORM_CLIP)
    dw = tf.math.minimum(dw, BBOX_XFORM_CLIP)

    anchor_ymin = anchors[..., 0:1]
    anchor_xmin = anchors[..., 1:2]
    anchor_ymax = anchors[..., 2:3]
    anchor_xmax = anchors[..., 3:4]
    anchor_h = anchor_ymax - anchor_ymin + 1.0
    anchor_w = anchor_xmax - anchor_xmin + 1.0
    anchor_yc = anchor_ymin + 0.5 * anchor_h
    anchor_xc = anchor_xmin + 0.5 * anchor_w

    decoded_boxes_yc = dy * anchor_h + anchor_yc
    decoded_boxes_xc = dx * anchor_w + anchor_xc
    decoded_boxes_h = tf.math.exp(dh) * anchor_h
    decoded_boxes_w = tf.math.exp(dw) * anchor_w

    decoded_boxes_ymin = decoded_boxes_yc - 0.5 * decoded_boxes_h
    decoded_boxes_xmin = decoded_boxes_xc - 0.5 * decoded_boxes_w
    decoded_boxes_ymax = decoded_boxes_ymin + decoded_boxes_h - 1.0
    decoded_boxes_xmax = decoded_boxes_xmin + decoded_boxes_w - 1.0

    decoded_boxes = tf.concat([
        decoded_boxes_ymin, decoded_boxes_xmin, decoded_boxes_ymax,
        decoded_boxes_xmax
    ],
                              axis=-1)
    return decoded_boxes
//  [python]
def mrcnn_orientation_loss_graph(target_orientation, target_class_ids, pred_orientation):
    """Loss for the orientation head of Mask RCNN.
    """
    # Loss
    # print("pred_orientation", pred_orientation.shape)
    target_class_ids = K.reshape(target_class_ids, (-1,))
    target_orientation = K.reshape(target_orientation, (-1, 1))
    pred_orientation = K.reshape(pred_orientation, (-1, K.int_shape(pred_orientation)[2], 1))
    # Only positive ROIs contribute to the loss. And only
    # the right class_id of each ROI. Get their indices.
    positive_roi_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_roi_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_roi_ix), tf.int64)
    indices = tf.stack([positive_roi_ix, positive_roi_class_ids], axis=1)

    # Gather the deltas (predicted and true) that contribute to loss
    target_orientation = tf.gather(target_orientation, positive_roi_ix)
    pred_orientation = tf.gather_nd(pred_orientation, indices)
    # print(target_orientation.shape)
    # print(pred_orientation.shape)
    # smooth L1 Loss
    # less_than_one = K.cast(K.less(diff, 1.0), "float32")
    # loss = (less_than_one * 0.5 * diff**2) + (1 - less_than_one) * (diff - 0.5)
    # loss = K.switch(tf.size(loss) > 0, K.mean(loss), tf.constant(0.0))
    # angle loss
    loss = tf.abs(tf.atan2(tf.sin(target_orientation - pred_orientation), tf.cos(target_orientation - pred_orientation)))
    loss = K.switch(tf.size(loss) > 0, K.mean(loss), tf.constant(0.0))
    
    return loss
//  [python]
def decode_boxes_lrtb(encoded_boxes_lrtb, anchors, weights=None):
  """Decode boxes.

  Args:
    encoded_boxes_lrtb: a tensor whose last dimension is 4 representing the
      coordinates of encoded boxes in left, right, top, bottom order.
    anchors: a tensor whose shape is the same as, or `broadcastable` to `boxes`,
      representing the coordinates of anchors in ymin, xmin, ymax, xmax order.
    weights: None or a list of four float numbers used to scale coordinates.

  Returns:
    decoded_boxes_lrtb: a tensor whose shape is the same as `boxes` representing
      the decoded box targets in lrtb (=left,right,top,bottom) format. The box
      decoded box coordinates represent the left, right, top, and bottom
      distances from an anchor location to the four borders of the matched
      groundtruth bounding box.
  """
  if encoded_boxes_lrtb.shape[-1] != 4:
    raise ValueError(
        'encoded_boxes_lrtb.shape[-1] is {:d}, but must be 4.'
        .format(encoded_boxes_lrtb.shape[-1]))

  with tf.name_scope('decode_boxes_lrtb'):
    encoded_boxes_lrtb = tf.to_float(encoded_boxes_lrtb, dtype=anchors.dtype)
    left = encoded_boxes_lrtb[..., 0:1]
    right = encoded_boxes_lrtb[..., 1:2]
    top = encoded_boxes_lrtb[..., 2:3]
    bottom = encoded_boxes_lrtb[..., 3:4]
    if weights:
      left /= weights[0]
      right /= weights[1]
      top /= weights[2]
      bottom /= weights[3]

    anchor_ymin = anchors[..., 0:1]
    anchor_xmin = anchors[..., 1:2]
    anchor_ymax = anchors[..., 2:3]
    anchor_xmax = anchors[..., 3:4]

    anchor_h = anchor_ymax - anchor_ymin
    anchor_w = anchor_xmax - anchor_xmin
    anchor_yc = anchor_ymin + 0.5 * anchor_h
    anchor_xc = anchor_xmin + 0.5 * anchor_w
    anchor_h += EPSILON
    anchor_w += EPSILON

    decoded_boxes_ymin = anchor_yc - top * anchor_h
    decoded_boxes_xmin = anchor_xc - left * anchor_w
    decoded_boxes_ymax = anchor_yc + bottom * anchor_h
    decoded_boxes_xmax = anchor_xc + right * anchor_w

    decoded_boxes_lrtb = tf.concat(
        [decoded_boxes_ymin, decoded_boxes_xmin,
         decoded_boxes_ymax, decoded_boxes_xmax],
        axis=-1)
    return decoded_boxes_lrtb
//  [python]
def decode_boxes_lrtb(encoded_boxes_lrtb, anchors, weights=None):
  """Decode boxes.

  Args:
    encoded_boxes_lrtb: a tensor whose last dimension is 4 representing the
      coordinates of encoded boxes in left, right, top, bottom order.
    anchors: a tensor whose shape is the same as, or `broadcastable` to `boxes`,
      representing the coordinates of anchors in ymin, xmin, ymax, xmax order.
    weights: None or a list of four float numbers used to scale coordinates.

  Returns:
    decoded_boxes_lrtb: a tensor whose shape is the same as `boxes` representing
      the decoded box targets in lrtb (=left,right,top,bottom) format. The box
      decoded box coordinates represent the left, right, top, and bottom
      distances from an anchor location to the four borders of the matched
      groundtruth bounding box.
  """
  if encoded_boxes_lrtb.shape[-1] != 4:
    raise ValueError(
        'encoded_boxes_lrtb.shape[-1] is {:d}, but must be 4.'
        .format(encoded_boxes_lrtb.shape[-1]))

  with tf.name_scope('decode_boxes_lrtb'):
    encoded_boxes_lrtb = tf.cast(encoded_boxes_lrtb, dtype=anchors.dtype)
    left = encoded_boxes_lrtb[..., 0:1]
    right = encoded_boxes_lrtb[..., 1:2]
    top = encoded_boxes_lrtb[..., 2:3]
    bottom = encoded_boxes_lrtb[..., 3:4]
    if weights:
      left /= weights[0]
      right /= weights[1]
      top /= weights[2]
      bottom /= weights[3]

    anchor_ymin = anchors[..., 0:1]
    anchor_xmin = anchors[..., 1:2]
    anchor_ymax = anchors[..., 2:3]
    anchor_xmax = anchors[..., 3:4]

    anchor_h = anchor_ymax - anchor_ymin
    anchor_w = anchor_xmax - anchor_xmin
    anchor_yc = anchor_ymin + 0.5 * anchor_h
    anchor_xc = anchor_xmin + 0.5 * anchor_w
    anchor_h += EPSILON
    anchor_w += EPSILON

    decoded_boxes_ymin = anchor_yc - top * anchor_h
    decoded_boxes_xmin = anchor_xc - left * anchor_w
    decoded_boxes_ymax = anchor_yc + bottom * anchor_h
    decoded_boxes_xmax = anchor_xc + right * anchor_w

    decoded_boxes_lrtb = tf.concat(
        [decoded_boxes_ymin, decoded_boxes_xmin,
         decoded_boxes_ymax, decoded_boxes_xmax],
        axis=-1)
    return decoded_boxes_lrtb
//  [python]
def mask_zeros_1(embed, lens, max_len, expand=True):
    mask = tf.sequence_mask(lens, max_len)
    mask = tf.to_float(mask)
    if(expand):
        mask = tf.expand_dims(mask, 2)
    embed = embed * mask
    return embed
//  [python]
def mask_zeros_1(embed, lens, max_len, expand=True):
    mask = tf.sequence_mask(lens, max_len)
    mask = tf.cast(mask, tf.float32)
    if(expand):
        mask = tf.expand_dims(mask, 2)
    embed = embed * mask
    return embed
//  [python]
def GumbelSoftmaxLogDensity(y, p, tau):
    # EPS = tf.constant(1e-10)
    k = tf.shape(y)[-1]
    k = t.cast(k, tf.float32)
    # y = y + EPS
    # y = tf.divide(y, tf.reduce_sum(y, -1, keep_dims=True))
    y = normalize_to_unit_sum(y)
    sum_p_over_y = tf.reduce_sum(tf.divide(p, tf.pow(y, tau)), -1)
    logp = tf.lgamma(k)
    logp = logp + (k - 1) * tf.log(tau)
    logp = logp - k * tf.log(sum_p_over_y)
    logp = logp + sum_p_over_y
    return logp
//  [python]
def GumbelSoftmaxLogDensity(y, p, tau):
    # EPS = tf.constant(1e-10)
    k = tf.shape(y)[-1]
    k = tf.cast(k, tf.float32)
    # y = y + EPS
    # y = tf.divide(y, tf.reduce_sum(y, -1, keep_dims=True))
    y = normalize_to_unit_sum(y)
    sum_p_over_y = tf.reduce_sum(tf.divide(p, tf.pow(y, tau)), -1)
    logp = tf.lgamma(k)
    logp = logp + (k - 1) * tf.log(tau)
    logp = logp - k * tf.log(sum_p_over_y)
    logp = logp + sum_p_over_y
    return logp
//  [python]
int_to_vocab, vocab_size = load_vocab(FLAGS.vocab_file)

vocab = lookup_ops.index_table_from_file(
  FLAGS.vocab_file, default_value=FLAGS.unk_id)
eos_id = tf.to_int32(vocab.lookup(tf.constant(FLAGS.eos)))
//  [python]
int_to_vocab, vocab_size = load_vocab(FLAGS.vocab_file)

vocab = lookup_ops.index_table_from_file(
  FLAGS.vocab_file, default_value=FLAGS.unk_id)
eos_id = tf.cast(vocab.lookup(tf.constant(FLAGS.eos)), tf.int32)