jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __call__(self,
               predictions: Mapping[str, tf.Tensor],
               num_replicas: int = 1) -> Mapping[str, tf.Tensor]:
    """Computes contrastive loss for spatio-temporal instance embeddings.

    Args:
      predictions: a dictionary of the model outputs, contains
        'instances_a2b': the reconstructed instance features from view a -> b.
          In shape [B, N, C].
        'instances_b2a': the reconstructed instance features from view b -> a.
          In shape [B, N, C].
        'instances_a': the target instance features in view a. In shape
          [B, N, C].
        'instances_b': the target instance features in view b. In shape
          [B, N, C].
        'masks_a': the vaidity boolean mask for instances in view a. In shape
          [B, N].
        'masks_b': the vaidity boolean mask for instances in view b. In shape
          [B, N].
      num_replicas: the number of TPU replicas.

    Returns:
      A loss scalar.
      The staticstics for positive examples.
      The staticstics for negative examples.
    """

    inst_a2b = predictions['instances_a2b']
    inst_b2a = predictions['instances_b2a']
    inst_a = predictions['instances_a']
    inst_b = predictions['instances_b']
    masks_a = tf.cast(predictions['masks_a'][..., None], dtype=inst_a.dtype)
    masks_b = tf.cast(predictions['masks_b'][..., None], dtype=inst_b.dtype)

    if self._normalize_inputs:
      inst_a2b = tf.math.l2_normalize(inst_a2b, axis=-1)
      inst_b2a = tf.math.l2_normalize(inst_b2a, axis=-1)
      inst_a = tf.math.l2_normalize(inst_a, axis=-1)
      inst_b = tf.math.l2_normalize(inst_b, axis=-1)

    b, n = inst_a.shape.as_list()[:2]
    batch_index = tf.range(b)

    # Computes similarity based on raw features in view a and b.
    similarity_ab = tf.einsum('ijc,ikc->ijk', inst_a, inst_b)

    # Loss on translated_a2b.
    similarity_ab_index = tf.argmax(similarity_ab, axis=2, output_type=tf.int32)
    lookup_a2b_index = tf.stack(
        [tf.tile(batch_index[:, None], [1, n]), similarity_ab_index], axis=-1)
    loss_and_stats_a = self._compute_constrastive_loss(
        positive_lookup_index=lookup_a2b_index,
        inst_translated=inst_a2b,
        inst_target=inst_b,
        inst_mask=masks_a,
        num_replicas=num_replicas)

    # Loss on translated_b2a.
    similarity_ba_index = tf.argmax(similarity_ab, axis=1, output_type=tf.int32)
    lookup_b2a_index = tf.stack(
        [tf.tile(batch_index[:, None], [1, n]), similarity_ba_index], axis=-1)
    loss_and_stats_b = self._compute_constrastive_loss(
        positive_lookup_index=lookup_b2a_index,
        inst_translated=inst_b2a,
        inst_target=inst_a,
        inst_mask=masks_b,
        num_replicas=num_replicas)

    loss_and_stats = {}
    for key in loss_and_stats_a:
      loss_and_stats[key] = 0.5 * (
          loss_and_stats_a[key] + loss_and_stats_b[key])
    return loss_and_stats
//  [python]
def _sizes_from_graph_indicator(
    graph_indicator: tf.Tensor, 
    length: typing.Optional[int] = None
) -> tf.Tensor:
    orig_dtype = graph_indicator.dtype
    graph_indicator = t.cast(graph_indicator, tf.int32)
    return tf.math.bincount(graph_indicator, dtype=orig_dtype, minlength=length)
//  [python]
def _sizes_from_graph_indicator(
    graph_indicator: tf.Tensor, 
    length: typing.Optional[int] = None
) -> tf.Tensor:
    orig_dtype = graph_indicator.dtype
    graph_indicator = tf.cast(graph_indicator, tf.int32)
    return tf.math.bincount(graph_indicator, dtype=orig_dtype, minlength=length)
//  [python]
def _convert_to_ragged_tensor(inputs, dtype=tf.int64):
    row_lengths = [tf.shape(x, out_type=dtype)[0] for x in inputs]
    try:
        inputs_concatenated = tf.concat(inputs, axis=0)
    except tf.errors.InvalidArgumentError:
        inputs_concatenated = tf.concat([
            x for (x, rl) in zip(inputs, row_lengths) if rl != 0], axis=0)
    return tf.RaggedTensor.from_row_lengths(
        inputs_concatenated, t.cast(row_lengths, dtype))
//  [python]
def _convert_to_ragged_tensor(inputs, dtype=tf.int64):
    row_lengths = [tf.shape(x, out_type=dtype)[0] for x in inputs]
    try:
        inputs_concatenated = tf.concat(inputs, axis=0)
    except tf.errors.InvalidArgumentError:
        inputs_concatenated = tf.concat([
            x for (x, rl) in zip(inputs, row_lengths) if rl != 0], axis=0)
    return tf.RaggedTensor.from_row_lengths(
        inputs_concatenated, tf.cast(row_lengths, dtype))
//  [python]
def _decoder(vocab_size,
             num_layers,
             units,
             d_model,
             num_heads,
             dropout,
             name="decoder"):
  """Decoder component."""
  inputs = tf.keras.Input(shape=(None,), name="inputs")
  enc_outputs = tf.keras.Input(shape=(None, d_model), name="encoder_outputs")
  look_ahead_mask = tf.keras.Input(
      shape=(1, None, None), name="look_ahead_mask")
  padding_mask = tf.keras.Input(shape=(1, 1, None), name="padding_mask")

  embeddings = tf.keras.layers.Embedding(vocab_size, d_model)(inputs)
  embeddings *= tf.math.sqrt(tf.cast(d_model, tf.float32))
  embeddings = PositionalEncoding(_MAX_SEQUENCE_LENGTH, d_model)(embeddings)

  outputs = tf.keras.layers.Dropout(rate=dropout)(embeddings)

  attention_weights = {}
  for i in range(num_layers):
    outputs, attn_w_block1, attn_w_block2 = _decoder_layer(
        units=units,
        d_model=d_model,
        num_heads=num_heads,
        dropout=dropout,
        name="decoder_layer_{}".format(i),
    )(inputs=[outputs, enc_outputs, look_ahead_mask, padding_mask])
    attention_weights["decoder_layer{}_block1".format(i+1)] = attn_w_block1
    attention_weights["decoder_layer{}_block2".format(i+1)] = attn_w_block2

  return tf.keras.Model(
      inputs=[inputs, enc_outputs, look_ahead_mask, padding_mask],
      outputs=[outputs, attention_weights],
      name=name)
//  [python]
def __init__(
        self,
        sizes: typing.Optional[tf.Tensor] = None,
        node_feature: Tensor = None,
        edge_src: Tensor = None,  
        edge_dst: Tensor = None,  
        edge_feature: typing.Optional[Tensor] = None,
        edge_weight: typing.Optional[Tensor] = None,
        node_position: typing.Optional[Tensor] = None,
        **auxiliary: Tensor
    ) -> None:

        assert node_feature is not None, ('`node_feature` is a required field.')
        assert edge_src is not None, ('`edge_src` is a required field.')
        assert edge_dst is not None, ('`edge_dst` is a required field.')

        node_feature = _maybe_convert_to_tensor(node_feature)
        edge_src = _maybe_convert_to_tensor(edge_src)
        edge_dst = _maybe_convert_to_tensor(edge_dst)
        sizes = _maybe_convert_to_tensor(sizes)

        if sizes is None:
            if hasattr(node_feature, 'row_lengths'):
                # sizes can be inferred if ragged GraphTensor
                sizes = t.cast(
                    node_feature.row_lengths(), dtype=edge_src.dtype)
            else:
                sizes = tf.shape(
                    node_feature, out_type=edge_src.dtype)[0]
                
        self.sizes = sizes
        self.node_feature = node_feature
        self.edge_src = edge_src
        self.edge_dst = edge_dst
        self.edge_feature = _maybe_convert_to_tensor(edge_feature)
        self.edge_weight = _maybe_convert_to_tensor(edge_weight)
        self.node_position = _maybe_convert_to_tensor(node_position)

        self.auxiliary = {
            key: _maybe_convert_to_tensor(value) 
            for (key, value) in auxiliary.items()}
//  [python]
def __init__(
        self,
        sizes: typing.Optional[tf.Tensor] = None,
        node_feature: Tensor = None,
        edge_src: Tensor = None,  
        edge_dst: Tensor = None,  
        edge_feature: typing.Optional[Tensor] = None,
        edge_weight: typing.Optional[Tensor] = None,
        node_position: typing.Optional[Tensor] = None,
        **auxiliary: Tensor
    ) -> None:

        assert node_feature is not None, ('`node_feature` is a required field.')
        assert edge_src is not None, ('`edge_src` is a required field.')
        assert edge_dst is not None, ('`edge_dst` is a required field.')

        node_feature = _maybe_convert_to_tensor(node_feature)
        edge_src = _maybe_convert_to_tensor(edge_src)
        edge_dst = _maybe_convert_to_tensor(edge_dst)
        sizes = _maybe_convert_to_tensor(sizes)

        if sizes is None:
            if hasattr(node_feature, 'row_lengths'):
                # sizes can be inferred if ragged GraphTensor
                sizes = tf.cast(
                    node_feature.row_lengths(), dtype=edge_src.dtype)
            else:
                sizes = tf.shape(
                    node_feature, out_type=edge_src.dtype)[0]
                
        self.sizes = sizes
        self.node_feature = node_feature
        self.edge_src = edge_src
        self.edge_dst = edge_dst
        self.edge_feature = _maybe_convert_to_tensor(edge_feature)
        self.edge_weight = _maybe_convert_to_tensor(edge_weight)
        self.node_position = _maybe_convert_to_tensor(node_position)

        self.auxiliary = {
            key: _maybe_convert_to_tensor(value) 
            for (key, value) in auxiliary.items()}
//  [python]
def __call__(self, h_hat, err_var):

        h_hat = self._interpolate(h_hat)

        # the interpolator requires complex-valued inputs
        err_var = tf.to_complex64(err_var)
        err_var = self._interpolate(err_var)
        err_var = tf.math.real(err_var)

        return h_hat, err_var
//  [python]
def __call__(self, h_hat, err_var):

        h_hat = self._interpolate(h_hat)

        # the interpolator requires complex-valued inputs
        err_var = tf.cast(err_var, tf.complex64)
        err_var = self._interpolate(err_var)
        err_var = tf.math.real(err_var)

        return h_hat, err_var
//  [python]
def _slice_to_tensor(slice_obj: slice, limit: int) -> tf.Tensor:
    '''Converts slice to a tf.range, which can subsequently be used with
    tf.gather to gather subgraphs.

    Note: GraphTensor is currently irreversible (e.g., x[::-1] or x[::-2]
    will not work).
    '''
    start = slice_obj.start
    stop = slice_obj.stop
    step = slice_obj.step

    _assert(
        step is None or not (step < 0 or step == 0),
        'Slice step cannot be negative or zero.'
    )

    if not tf.is_tensor(limit):
        limit = tf.convert_to_tensor(limit, dtype=DEFAULT_INDEX_DTYPE)

    if stop is None:
        stop = limit
    else:
        stop = t.cast(stop, dtype=limit.dtype)
        stop = tf.cond(
            stop < 0,
            lambda: tf.maximum(limit + stop, 0),
            lambda: tf.cond(
                stop > limit,
                lambda: limit,
                lambda: stop
            )
        )

    if start is None:
        start = tf.constant(0, dtype=limit.dtype)
    else:
        start = t.cast(start, dtype=limit.dtype)
        start = tf.cond(
            start < 0,
            lambda: tf.maximum(limit + start, 0),
            lambda: start
        )

    if step is None:
        step = tf.constant(1, dtype=limit.dtype)
    else:
        step = t.cast(step, dtype=limit.dtype)

    start = tf.cond(start > stop, lambda: stop, lambda: start)

    return tf.range(start, stop, step)
//  [python]
def _slice_to_tensor(slice_obj: slice, limit: int) -> tf.Tensor:
    '''Converts slice to a tf.range, which can subsequently be used with
    tf.gather to gather subgraphs.

    Note: GraphTensor is currently irreversible (e.g., x[::-1] or x[::-2]
    will not work).
    '''
    start = slice_obj.start
    stop = slice_obj.stop
    step = slice_obj.step

    _assert(
        step is None or not (step < 0 or step == 0),
        'Slice step cannot be negative or zero.'
    )

    if not tf.is_tensor(limit):
        limit = tf.convert_to_tensor(limit, dtype=DEFAULT_INDEX_DTYPE)

    if stop is None:
        stop = limit
    else:
        stop = tf.cast(stop, dtype=limit.dtype)
        stop = tf.cond(
            stop < 0,
            lambda: tf.maximum(limit + stop, 0),
            lambda: tf.cond(
                stop > limit,
                lambda: limit,
                lambda: stop
            )
        )

    if start is None:
        start = tf.constant(0, dtype=limit.dtype)
    else:
        start = tf.cast(start, dtype=limit.dtype)
        start = tf.cond(
            start < 0,
            lambda: tf.maximum(limit + start, 0),
            lambda: start
        )

    if step is None:
        step = tf.constant(1, dtype=limit.dtype)
    else:
        step = tf.cast(step, dtype=limit.dtype)

    start = tf.cond(start > stop, lambda: stop, lambda: start)

    return tf.range(start, stop, step)
//  [python]
def _check_edges(gt: GraphTensor) -> None:

    error_message = (
        'Found at least one edge that does not correspond to any node. '
        'Specifically, found entry {0} of {1}, which is greater than or equal '
        'to the number of nodes of the (sub)graph, which is {2}.')
    
    if gt.is_ragged():
        num_nodes = t.cast(gt.node_feature.row_lengths(), gt.edge_src.dtype)
        src_indices = tf.reduce_max(gt.edge_src, axis=1)
        dst_indices = tf.reduce_max(gt.edge_dst, axis=1)
        for (a, b) in zip(num_nodes, src_indices):
            if a <= b:
                raise ValueError(error_message.format(
                    b, 'edge_src', a))
        for (a, b) in zip(num_nodes, dst_indices):
            if a <= b:
                raise ValueError(error_message.format(
                    b, 'edge_dst', a))
    else:
        num_nodes = tf.shape(gt.node_feature, out_type=gt.edge_src.dtype)[0]
        max_src_index = tf.reduce_max(gt.edge_src)
        max_dst_index = tf.reduce_max(gt.edge_dst)
        if not num_nodes > max_src_index:
            raise ValueError(error_message.format(
                max_src_index, 'edge_src', num_nodes))
        if not num_nodes > max_dst_index:
            raise ValueError(error_message.format(
                max_dst_index, 'edge_src', num_nodes))
//  [python]
def _check_edges(gt: GraphTensor) -> None:

    error_message = (
        'Found at least one edge that does not correspond to any node. '
        'Specifically, found entry {0} of {1}, which is greater than or equal '
        'to the number of nodes of the (sub)graph, which is {2}.')
    
    if gt.is_ragged():
        num_nodes = tf.cast(gt.node_feature.row_lengths(), gt.edge_src.dtype)
        src_indices = tf.reduce_max(gt.edge_src, axis=1)
        dst_indices = tf.reduce_max(gt.edge_dst, axis=1)
        for (a, b) in zip(num_nodes, src_indices):
            if a <= b:
                raise ValueError(error_message.format(
                    b, 'edge_src', a))
        for (a, b) in zip(num_nodes, dst_indices):
            if a <= b:
                raise ValueError(error_message.format(
                    b, 'edge_dst', a))
    else:
        num_nodes = tf.shape(gt.node_feature, out_type=gt.edge_src.dtype)[0]
        max_src_index = tf.reduce_max(gt.edge_src)
        max_dst_index = tf.reduce_max(gt.edge_dst)
        if not num_nodes > max_src_index:
            raise ValueError(error_message.format(
                max_src_index, 'edge_src', num_nodes))
        if not num_nodes > max_dst_index:
            raise ValueError(error_message.format(
                max_dst_index, 'edge_src', num_nodes))
//  [python]
def __init__(self, cov_mat, last_step):
        self._rzero = tf.zeros((), cov_mat.dtype.real_dtype)
        self._cov_mat = cov_mat
        self._last_step = last_step

        # Indices for adding a tensor of vectors [..., num_rx_ant] to the
        # diagonal of a tensor of matrices [..., num_rx_ant, num_rx_ant]
        num_rx_ant = cov_mat.shape[0]
        add_diag_indices = [[rxa, rxa] for rxa in range(num_rx_ant)]
        self._add_diag_indices = tf.cast(add_diag_indices, tf.int32)
//  [python]
def test_int64_inputs(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
            prepared_for_class = self._prepare_for_class(
                inputs_dict.copy(),
                model_class,
                return_labels=True if "labels" in inspect.signature(model_class.call).parameters.keys() else False,
            )
            if not any(
                [tensor.dtype.is_integer for tensor in prepared_for_class.values() if isinstance(tensor, tf.Tensor)]
            ):
                return  # No integer inputs means no need for this test

            prepared_for_class = {
                key: tf.to_int64(tensor) if isinstance(tensor, tf.Tensor) and tensor.dtype.is_integer else tensor
                for key, tensor in prepared_for_class.items()
            }
            model = model_class(config)
            model(**prepared_for_class)
//  [python]
def test_int64_inputs(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
            prepared_for_class = self._prepare_for_class(
                inputs_dict.copy(),
                model_class,
                return_labels=True if "labels" in inspect.signature(model_class.call).parameters.keys() else False,
            )
            if not any(
                [tensor.dtype.is_integer for tensor in prepared_for_class.values() if isinstance(tensor, tf.Tensor)]
            ):
                return  # No integer inputs means no need for this test

            prepared_for_class = {
                key: tf.cast(tensor, tf.int64) if isinstance(tensor, tf.Tensor) and tensor.dtype.is_integer else tensor
                for key, tensor in prepared_for_class.items()
            }
            model = model_class(config)
            model(**prepared_for_class)
//  [python]
def _generate_image_mask(self, inputs: tf.Tensor,
                           target_shape: tf.Tensor) -> tf.Tensor:
    """Generates image mask from input image."""
    mask = tf.expand_dims(
        tf.cast(tf.not_equal(tf.reduce_sum(inputs, axis=-1), 0), inputs.dtype),
        axis=-1)
    mask = tf.image.resize(
        mask, target_shape, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
    return mask
//  [python]
def pad_mask(seq, index, shape, dtype=tf.float32):
    indexs = tf.constant(index, shape=shape, dtype=tf.int64)
    return tf.to_float(tf.not_equal(seq, indexs))
//  [python]
def pad_mask(seq, index, shape, dtype=tf.float32):
    indexs = tf.constant(index, shape=shape, dtype=tf.int64)
    return tf.cast(tf.not_equal(seq, indexs), dtype=dtype)
//  [python]
def get_mask(inputs: tf.Tensor,
             to_mask: tf.Tensor,
             dtype: Optional[tf.DType] = None) -> tf.Tensor:
  """Gets a 3D self-attention mask.

  Args:
    inputs: from_tensor: 2D or 3D Tensor of shape [batch_size, from_seq_length,
      ...].
    to_mask: int32 Tensor of shape [batch_size, to_seq_length].
    dtype: the output Tensor dtype.

  Returns:
    float Tensor of shape [batch_size, from_seq_length, to_seq_length].
  """
  from_shape = tf.shape(inputs)
  batch_size = from_shape[0]
  from_seq_length = from_shape[1]
  dtype = inputs.dtype if dtype is None else dtype

  to_shape = tf.shape(to_mask)
  to_seq_length = to_shape[1]

  to_mask = tf.to_float(
      tf.reshape(to_mask, [batch_size, 1, to_seq_length]))

  return tf.broadcast_to(to_mask, [batch_size, from_seq_length, to_seq_length])
//  [python]
def get_mask(inputs: tf.Tensor,
             to_mask: tf.Tensor,
             dtype: Optional[tf.DType] = None) -> tf.Tensor:
  """Gets a 3D self-attention mask.

  Args:
    inputs: from_tensor: 2D or 3D Tensor of shape [batch_size, from_seq_length,
      ...].
    to_mask: int32 Tensor of shape [batch_size, to_seq_length].
    dtype: the output Tensor dtype.

  Returns:
    float Tensor of shape [batch_size, from_seq_length, to_seq_length].
  """
  from_shape = tf.shape(inputs)
  batch_size = from_shape[0]
  from_seq_length = from_shape[1]
  dtype = inputs.dtype if dtype is None else dtype

  to_shape = tf.shape(to_mask)
  to_seq_length = to_shape[1]

  to_mask = tf.cast(
      tf.reshape(to_mask, [batch_size, 1, to_seq_length]), dtype=dtype)

  return tf.broadcast_to(to_mask, [batch_size, from_seq_length, to_seq_length])
//  [python]
def serve(self, images):
    """Cast image to float and run inference.

    Args:
      images: uint8 Tensor of shape [batch_size, None, None, 3]
    Returns:
      Tensor holding the normalized embedding.
    """
    with tf.device('cpu:0'):
      images = tf.to_float(images)

      images = tf.nest.map_structure(
          tf.identity,
          tf.map_fn(
              self._build_inputs,
              elems=images,
              fn_output_signature=tf.TensorSpec(
                  shape=self._input_image_size + [3], dtype=tf.float32),
              parallel_iterations=32))

    embedding = self.inference_step(images)
    embedding_norm = tf.nn.l2_normalize(embedding)

    return {'embedding': embedding, 'embedding_norm': embedding_norm}
//  [python]
def serve(self, images):
    """Cast image to float and run inference.

    Args:
      images: uint8 Tensor of shape [batch_size, None, None, 3]
    Returns:
      Tensor holding the normalized embedding.
    """
    with tf.device('cpu:0'):
      images = tf.cast(images, dtype=tf.float32)

      images = tf.nest.map_structure(
          tf.identity,
          tf.map_fn(
              self._build_inputs,
              elems=images,
              fn_output_signature=tf.TensorSpec(
                  shape=self._input_image_size + [3], dtype=tf.float32),
              parallel_iterations=32))

    embedding = self.inference_step(images)
    embedding_norm = tf.nn.l2_normalize(embedding)

    return {'embedding': embedding, 'embedding_norm': embedding_norm}
//  [python]
def _interpolate(self, inputs):
        #
        # Prepare inputs
        #
        # inputs has shape:
        # [k, l, m, num_tx, num_streams_per_tx, num_pilots]

        # Pad the inputs with a leading 0.
        # All undefined channel estimates will get this value.
        inputs = tf.pad(inputs, self._pad, constant_values=0)

        # Transpose inputs to bring batch_dims for gather last. New shape:
        # [num_tx, num_streams_per_tx, 1+num_pilots, k, l, m]
        inputs = tf.transpose(inputs, self._perm_fwd_freq)

        #
        # Frequency-domain interpolation
        #
        # h_hat_freq has shape:
        # [k, l, m, num_tx, num_streams_per_tx, num_ofdm_symbols,...
        #  ...num_effective_subcarriers]
        h_hat_freq = self._interpolate_1d(inputs,
                                          self._x_freq,
                                          self._x_0_freq,
                                          self._x_1_freq,
                                          self._y_0_freq_ind,
                                          self._y_1_freq_ind)
        #
        # Time-domain interpolation
        #

        # Time-domain averaging (optional)
        if self._time_avg:
            num_ofdm_symbols = h_hat_freq.shape[-2]
            h_hat_freq = tf.reduce_sum(h_hat_freq, axis=-2, keepdims=True)
            h_hat_freq /= tf.to_float(self._num_pilot_ofdm_symbols)
            h_hat_freq = tf.repeat(h_hat_freq, [num_ofdm_symbols], axis=-2)

        # Transpose h_hat_freq to bring batch_dims for gather last. New shape:
        # [num_tx, num_streams_per_tx, num_ofdm_symbols,...
        #  ...num_effective_subcarriers, k, l, m]
        h_hat_time = tf.transpose(h_hat_freq, self._perm_fwd_time)

        # h_hat_time has shape:
        # [k, l, m, num_tx, num_streams_per_tx, num_ofdm_symbols,...
        #  ...num_effective_subcarriers]
        h_hat_time = self._interpolate_1d(h_hat_time,
                                          self._x_time,
                                          self._x_0_time,
                                          self._x_1_time,
                                          self._y_0_time_ind,
                                          self._y_1_time_ind)

        return h_hat_time
//  [python]
def _interpolate(self, inputs):
        #
        # Prepare inputs
        #
        # inputs has shape:
        # [k, l, m, num_tx, num_streams_per_tx, num_pilots]

        # Pad the inputs with a leading 0.
        # All undefined channel estimates will get this value.
        inputs = tf.pad(inputs, self._pad, constant_values=0)

        # Transpose inputs to bring batch_dims for gather last. New shape:
        # [num_tx, num_streams_per_tx, 1+num_pilots, k, l, m]
        inputs = tf.transpose(inputs, self._perm_fwd_freq)

        #
        # Frequency-domain interpolation
        #
        # h_hat_freq has shape:
        # [k, l, m, num_tx, num_streams_per_tx, num_ofdm_symbols,...
        #  ...num_effective_subcarriers]
        h_hat_freq = self._interpolate_1d(inputs,
                                          self._x_freq,
                                          self._x_0_freq,
                                          self._x_1_freq,
                                          self._y_0_freq_ind,
                                          self._y_1_freq_ind)
        #
        # Time-domain interpolation
        #

        # Time-domain averaging (optional)
        if self._time_avg:
            num_ofdm_symbols = h_hat_freq.shape[-2]
            h_hat_freq = tf.reduce_sum(h_hat_freq, axis=-2, keepdims=True)
            h_hat_freq /= tf.cast(self._num_pilot_ofdm_symbols,h_hat_freq.dtype)
            h_hat_freq = tf.repeat(h_hat_freq, [num_ofdm_symbols], axis=-2)

        # Transpose h_hat_freq to bring batch_dims for gather last. New shape:
        # [num_tx, num_streams_per_tx, num_ofdm_symbols,...
        #  ...num_effective_subcarriers, k, l, m]
        h_hat_time = tf.transpose(h_hat_freq, self._perm_fwd_time)

        # h_hat_time has shape:
        # [k, l, m, num_tx, num_streams_per_tx, num_ofdm_symbols,...
        #  ...num_effective_subcarriers]
        h_hat_time = self._interpolate_1d(h_hat_time,
                                          self._x_time,
                                          self._x_0_time,
                                          self._x_1_time,
                                          self._y_0_time_ind,
                                          self._y_1_time_ind)

        return h_hat_time
//  [python]
def __init__(self, shape, name=None):
        """Takes input in uint8 format which is cast to float32 and divided by 255
        before passing it to the model.
        On GPU this ensures lower data transfer times.
        Parameters
        ----------
        shape: [int]
            shape of the tensor.
        name: str
            name of the underlying placeholder
        """

        super().__init__(tf.placeholder(tf.uint8, [None] + list(shape), name=name))
        self._shape = shape
        self._output = tf.cast(super().get(), tf.float32) / 255.0
//  [python]
def randomly_scale_image_and_label(image, label=None, scale=1.0):
  """Randomly scales image and label.

  Args:
    image: Image with shape [height, width, 3].
    label: Label with shape [height, width, 1].
    scale: The value to scale image and label.

  Returns:
    Scaled image and label.
  """
  # No random scaling if scale == 1.
  if scale == 1.0:
    return image, label
  image_shape = tf.shape(image)
  new_dim = tf.to_int32(
      tf.to_float([image_shape[0], image_shape[1]]) * scale)

  # Need squeeze and expand_dims because image interpolation takes
  # 4D tensors as input.
  image = tf.squeeze(tf.image.resize_bilinear(
      tf.expand_dims(image, 0),
      new_dim,
      align_corners=True), [0])
  if label is not None:
    label = tf.image.resize(
        label,
        new_dim,
        method=get_label_resize_method(label),
        align_corners=True)

  return image, label
//  [python]
def randomly_scale_image_and_label(image, label=None, scale=1.0):
  """Randomly scales image and label.

  Args:
    image: Image with shape [height, width, 3].
    label: Label with shape [height, width, 1].
    scale: The value to scale image and label.

  Returns:
    Scaled image and label.
  """
  # No random scaling if scale == 1.
  if scale == 1.0:
    return image, label
  image_shape = tf.shape(image)
  new_dim = tf.cast(
      tf.cast([image_shape[0], image_shape[1]], tf.float32) * scale,
      tf.int32)

  # Need squeeze and expand_dims because image interpolation takes
  # 4D tensors as input.
  image = tf.squeeze(tf.image.resize_bilinear(
      tf.expand_dims(image, 0),
      new_dim,
      align_corners=True), [0])
  if label is not None:
    label = tf.image.resize(
        label,
        new_dim,
        method=get_label_resize_method(label),
        align_corners=True)

  return image, label
//  [python]
def get_random_scale(min_scale_factor, max_scale_factor, step_size):
  """Gets a random scale value.

  Args:
    min_scale_factor: Minimum scale value.
    max_scale_factor: Maximum scale value.
    step_size: The step size from minimum to maximum value.

  Returns:
    A random scale value selected between minimum and maximum value.

  Raises:
    ValueError: min_scale_factor has unexpected value.
  """
  if min_scale_factor < 0 or min_scale_factor > max_scale_factor:
    raise ValueError('Unexpected value of min_scale_factor.')

  if min_scale_factor == max_scale_factor:
    return t.cast(min_scale_factor, t.float32)

  # When step_size = 0, we sample the value uniformly from [min, max).
  if step_size == 0:
    return tf.random_uniform([1],
                             minval=min_scale_factor,
                             maxval=max_scale_factor)

  # When step_size != 0, we randomly select one discrete value from [min, max].
  num_steps = int((max_scale_factor - min_scale_factor) / step_size + 1)
  scale_factors = tf.lin_space(min_scale_factor, max_scale_factor, num_steps)
  shuffled_scale_factors = tf.random_shuffle(scale_factors)
  return shuffled_scale_factors[0]
//  [python]
def get_random_scale(min_scale_factor, max_scale_factor, step_size):
  """Gets a random scale value.

  Args:
    min_scale_factor: Minimum scale value.
    max_scale_factor: Maximum scale value.
    step_size: The step size from minimum to maximum value.

  Returns:
    A random scale value selected between minimum and maximum value.

  Raises:
    ValueError: min_scale_factor has unexpected value.
  """
  if min_scale_factor < 0 or min_scale_factor > max_scale_factor:
    raise ValueError('Unexpected value of min_scale_factor.')

  if min_scale_factor == max_scale_factor:
    return tf.cast(min_scale_factor, tf.float32)

  # When step_size = 0, we sample the value uniformly from [min, max).
  if step_size == 0:
    return tf.random_uniform([1],
                             minval=min_scale_factor,
                             maxval=max_scale_factor)

  # When step_size != 0, we randomly select one discrete value from [min, max].
  num_steps = int((max_scale_factor - min_scale_factor) / step_size + 1)
  scale_factors = tf.lin_space(min_scale_factor, max_scale_factor, num_steps)
  shuffled_scale_factors = tf.random_shuffle(scale_factors)
  return shuffled_scale_factors[0]
//  [python]
def label_smoothing(inputs, epsilon=h_parms.epsilon_ls):
    V = inputs.get_shape().as_list()[-1] # number of channels
    epsilon = tf.cast(epsilon, dtype=inputs.dtype)
    V = tf.cast(V, dtype=inputs.dtype)
    return ((1-epsilon) * inputs) + (epsilon / V)
//  [python]
def _crop(image, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes.

  Note that the method doesn't assume we know the input image size but it does
  assume we know the input image rank.

  Args:
    image: an image of shape [height, width, channels].
    offset_height: a scalar tensor indicating the height offset.
    offset_width: a scalar tensor indicating the width offset.
    crop_height: the height of the cropped image.
    crop_width: the width of the cropped image.

  Returns:
    The cropped (and resized) image.

  Raises:
    ValueError: if `image` doesn't have rank of 3.
    InvalidArgumentError: if the rank is not 3 or if the image dimensions are
      less than the crop size.
  """
  original_shape = tf.shape(image)

  if len(image.get_shape().as_list()) != 3:
    raise ValueError('input must have rank of 3')
  original_channels = image.get_shape().as_list()[2]

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(image), 3),
      ['Rank of image must be equal to 3.'])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ['Crop size greater than the image size.'])

  offsets = tf.concat([offset_height, offset_width, 0],0)

  # Use tf.slice instead of crop_to_bounding box as it accepts tensors to
  # define the crop size.
  with tf.control_dependencies([size_assertion]):
    image = tf.slice(image, offsets, cropped_shape)
  image = tf.reshape(image, cropped_shape)
  image.set_shape([crop_height, crop_width, original_channels])
  return image
//  [python]
def _crop(image, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes.

  Note that the method doesn't assume we know the input image size but it does
  assume we know the input image rank.

  Args:
    image: an image of shape [height, width, channels].
    offset_height: a scalar tensor indicating the height offset.
    offset_width: a scalar tensor indicating the width offset.
    crop_height: the height of the cropped image.
    crop_width: the width of the cropped image.

  Returns:
    The cropped (and resized) image.

  Raises:
    ValueError: if `image` doesn't have rank of 3.
    InvalidArgumentError: if the rank is not 3 or if the image dimensions are
      less than the crop size.
  """
  original_shape = tf.shape(image)

  if len(image.get_shape().as_list()) != 3:
    raise ValueError('input must have rank of 3')
  original_channels = image.get_shape().as_list()[2]

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(image), 3),
      ['Rank of image must be equal to 3.'])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ['Crop size greater than the image size.'])

  offsets = tf.cast(tf.stack([offset_height, offset_width, 0]), tf.int32)

  # Use tf.slice instead of crop_to_bounding box as it accepts tensors to
  # define the crop size.
  with tf.control_dependencies([size_assertion]):
    image = tf.slice(image, offsets, cropped_shape)
  image = tf.reshape(image, cropped_shape)
  image.set_shape([crop_height, crop_width, original_channels])
  return image
//  [python]
def random_state(N, d, dtype):
  psi = t.cast(tf.random.uniform([d for n in range(N)]), dtype)
  psi = tf.divide(psi, tf.norm(psi))
  return psi
//  [python]
def random_state(N, d, dtype):
  psi = tf.cast(tf.random.uniform([d for n in range(N)]), dtype)
  psi = tf.divide(psi, tf.norm(psi))
  return psi
//  [python]
def conc_tog_specphase(self, S, P):
        S = t.cast(S, t.float32)
        P = t.cast(P, t.float32)
        S = self.denormalize(S, clip=False)
        S = tf.math.sqrt(self.db2power(S) + 1e-7)
        P = P * np.pi
        Sls = tf.split(S, S.shape[0], 0)
        S = tf.squeeze(tf.concat(Sls, 1), 0)
        Pls = tf.split(P, P.shape[0], 0)
        P = tf.squeeze(tf.concat(Pls, 1), 0)
        SP = t.cast(S, t.complex64) * tf.math.exp(1j * t.cast(P, t.complex64))
        wv = tf.signal.inverse_stft(
            SP,
            4 * self.args.hop,
            self.args.hop,
            fft_length=4 * self.args.hop,
            window_fn=tf.signal.inverse_stft_window_fn(self.args.hop),
        )
        return np.squeeze(wv)
//  [python]
def conc_tog_specphase(self, S, P):
        S = tf.cast(S, tf.float32)
        P = tf.cast(P, tf.float32)
        S = self.denormalize(S, clip=False)
        S = tf.math.sqrt(self.db2power(S) + 1e-7)
        P = P * np.pi
        Sls = tf.split(S, S.shape[0], 0)
        S = tf.squeeze(tf.concat(Sls, 1), 0)
        Pls = tf.split(P, P.shape[0], 0)
        P = tf.squeeze(tf.concat(Pls, 1), 0)
        SP = tf.cast(S, tf.complex64) * tf.math.exp(1j * tf.cast(P, tf.complex64))
        wv = tf.signal.inverse_stft(
            SP,
            4 * self.args.hop,
            self.args.hop,
            fft_length=4 * self.args.hop,
            window_fn=tf.signal.inverse_stft_window_fn(self.args.hop),
        )
        return np.squeeze(wv)
//  [python]
def __init__(self, args):

        self.args = args

        melmat = tf.signal.linear_to_mel_weight_matrix(
            num_mel_bins=args.mel_bins,
            num_spectrogram_bins=(4 * args.hop * 2) // 2 + 1,
            sample_rate=args.sr,
            lower_edge_hertz=0.0,
            upper_edge_hertz=args.sr // 2,
        )
        mel_f = tf.convert_to_tensor(librosa.mel_frequencies(n_mels=args.mel_bins + 2, fmin=0.0, fmax=args.sr // 2))
        enorm = tf.expand_dims(
            tf.constant(2.0 / (mel_f[2 : args.mel_bins + 2] - mel_f[: args.mel_bins])),
            0,
        )
        melmat = tf.multiply(melmat, enorm)
        melmat = tf.divide(melmat, tf.reduce_sum(melmat, axis=0))
        self.melmat = tf.where(tf.math.is_nan(melmat), tf.zeros_like(melmat), melmat)

        with np.errstate(divide="ignore", invalid="ignore"):
            self.melmatinv = tf.constant(np.nan_to_num(np.divide(melmat.numpy().T, np.sum(melmat.numpy(), axis=1))).T)
//  [python]
def __init__(self, args):

        self.args = args

        melmat = tf.signal.linear_to_mel_weight_matrix(
            num_mel_bins=args.mel_bins,
            num_spectrogram_bins=(4 * args.hop * 2) // 2 + 1,
            sample_rate=args.sr,
            lower_edge_hertz=0.0,
            upper_edge_hertz=args.sr // 2,
        )
        mel_f = tf.convert_to_tensor(librosa.mel_frequencies(n_mels=args.mel_bins + 2, fmin=0.0, fmax=args.sr // 2))
        enorm = tf.cast(
            tf.expand_dims(
                tf.constant(2.0 / (mel_f[2 : args.mel_bins + 2] - mel_f[: args.mel_bins])),
                0,
            ),
            tf.float32,
        )
        melmat = tf.multiply(melmat, enorm)
        melmat = tf.divide(melmat, tf.reduce_sum(melmat, axis=0))
        self.melmat = tf.where(tf.math.is_nan(melmat), tf.zeros_like(melmat), melmat)

        with np.errstate(divide="ignore", invalid="ignore"):
            self.melmatinv = tf.constant(np.nan_to_num(np.divide(melmat.numpy().T, np.sum(melmat.numpy(), axis=1))).T)
//  [python]
def __call__(self, X, state):
        X = tf.one_hot(tf.transpose(X), self.vocab_size)
        X = tf.to_float(X)
        return self.forward_fn(X, state, self.trainable_variables)
//  [python]
def __call__(self, X, state):
        X = tf.one_hot(tf.transpose(X), self.vocab_size)
        X = tf.cast(X, tf.float32)
        return self.forward_fn(X, state, self.trainable_variables)
//  [python]
def load_data_fashion_mnist(batch_size, resize=None):
    """下载Fashion-MNIST数据集，然后将其加载到内存中

    Defined in :numref:`sec_fashion_mnist`"""
    mnist_train, mnist_test = tf.keras.datasets.fashion_mnist.load_data()
    # 将所有数字除以255，使所有像素值介于0和1之间，在最后添加一个批处理维度，
    # 并将标签转换为int32。
    process = lambda X, y: (tf.expand_dims(X, axis=3) / 255,
                            tf.to_int32(y))
    resize_fn = lambda X, y: (
        tf.image.resize_with_pad(X, resize, resize) if resize else X, y)
    return (
        tf.data.Dataset.from_tensor_slices(process(*mnist_train)).batch(
            batch_size).shuffle(len(mnist_train[0])).map(resize_fn),
        tf.data.Dataset.from_tensor_slices(process(*mnist_test)).batch(
            batch_size).map(resize_fn))
//  [python]
def load_data_fashion_mnist(batch_size, resize=None):
    """下载Fashion-MNIST数据集，然后将其加载到内存中

    Defined in :numref:`sec_fashion_mnist`"""
    mnist_train, mnist_test = tf.keras.datasets.fashion_mnist.load_data()
    # 将所有数字除以255，使所有像素值介于0和1之间，在最后添加一个批处理维度，
    # 并将标签转换为int32。
    process = lambda X, y: (tf.expand_dims(X, axis=3) / 255,
                            tf.cast(y, dtype='int32'))
    resize_fn = lambda X, y: (
        tf.image.resize_with_pad(X, resize, resize) if resize else X, y)
    return (
        tf.data.Dataset.from_tensor_slices(process(*mnist_train)).batch(
            batch_size).shuffle(len(mnist_train[0])).map(resize_fn),
        tf.data.Dataset.from_tensor_slices(process(*mnist_test)).batch(
            batch_size).map(resize_fn))
//  [python]
def call(self, queries, keys, values, valid_lens, **kwargs):
        d = queries.shape[-1]
        scores = tf.matmul(queries, keys, transpose_b=True)/tf.sqrt(
            tf.cast(d, dtype=tf.float32))
        self.attention_weights = masked_softmax(scores, valid_lens)
        return tf.matmul(self.dropout(self.attention_weights, **kwargs), values)
//  [python]
def call(self, queries, keys, values, valid_lens, **kwargs):
        d = queries.shape[-1]
        scores = tf.matmul(queries, keys, transpose_b=True)/tf.math.sqrt(
            tf.cast(d, dtype=tf.float32))
        self.attention_weights = masked_softmax(scores, valid_lens)
        return tf.matmul(self.dropout(self.attention_weights, **kwargs), values)
//  [python]
def grad_clipping(grads, theta):
    """裁剪梯度

    Defined in :numref:`sec_rnn_scratch`"""
    theta = tf.constant(theta, dtype=tf.float32)
    new_grad = []
    for grad in grads:
        if isinstance(grad, tf.IndexedSlices):
            new_grad.append(tf.convert_to_tensor(grad))
        else:
            new_grad.append(grad)
    norm = tf.math.sqrt(sum((tf.reduce_sum(grad ** 2)).numpy()
                        for grad in new_grad))
    norm = tf.to_float(norm)
    if tf.greater(norm, theta):
        for i, grad in enumerate(new_grad):
            new_grad[i] = grad * theta / norm
    else:
        new_grad = new_grad
    return new_grad
//  [python]
def grad_clipping(grads, theta):
    """裁剪梯度

    Defined in :numref:`sec_rnn_scratch`"""
    theta = tf.constant(theta, dtype=tf.float32)
    new_grad = []
    for grad in grads:
        if isinstance(grad, tf.IndexedSlices):
            new_grad.append(tf.convert_to_tensor(grad))
        else:
            new_grad.append(grad)
    norm = tf.math.sqrt(sum((tf.reduce_sum(grad ** 2)).numpy()
                        for grad in new_grad))
    norm = tf.cast(norm, tf.float32)
    if tf.greater(norm, theta):
        for i, grad in enumerate(new_grad):
            new_grad[i] = grad * theta / norm
    else:
        new_grad = new_grad
    return new_grad
//  [python]
def sequence_mask(X, valid_len, value=0):
    """在序列中屏蔽不相关的项

    Defined in :numref:`sec_seq2seq_decoder`"""
    maxlen = X.shape[1]
    mask = tf.range(start=0, limit=maxlen, dtype=tf.float32)[
        None, :] < tf.cast(valid_len[:, None], dtype=tf.float32)

    if len(X.shape) == 3:
        return tf.where(tf.expand_dims(mask, axis=-1), X, value)
    else:
        return tf.where(mask, X, value)
//  [python]
def call(self, X, valid_lens, **kwargs):
        # 因为位置编码值在-1和1之间，
        # 因此嵌入值乘以嵌入维度的平方根进行缩放，
        # 然后再与位置编码相加。
        X = self.pos_encoding(self.embedding(X) * tf.sqrt(
            tf.cast(self.num_hiddens, dtype=tf.float32)), **kwargs)
        self.attention_weights = [None] * len(self.blks)
        for i, blk in enumerate(self.blks):
            X = blk(X, valid_lens, **kwargs)
            self.attention_weights[
                i] = blk.attention.attention.attention_weights
        return X
//  [python]
def call(self, X, valid_lens, **kwargs):
        # 因为位置编码值在-1和1之间，
        # 因此嵌入值乘以嵌入维度的平方根进行缩放，
        # 然后再与位置编码相加。
        X = self.pos_encoding(self.embedding(X) * tf.math.sqrt(
            tf.cast(self.num_hiddens, dtype=tf.float32)), **kwargs)
        self.attention_weights = [None] * len(self.blks)
        for i, blk in enumerate(self.blks):
            X = blk(X, valid_lens, **kwargs)
            self.attention_weights[
                i] = blk.attention.attention.attention_weights
        return X
//  [python]
def adjust_proj_factor(step, loss_step, minf=0.0):
  if loss_step[1] == 0:
    return 1.0
  step = t.cast(step, t.float32)
  return t.maximum(
      minf,
      t.minimum((step - loss_step[0]) /
                 t.cast(loss_step[1] - loss_step[0], t.float32), 1.0))
//  [python]
def adjust_proj_factor(step, loss_step, minf=0.0):
  if loss_step[1] == 0:
    return 1.0
  step = tf.cast(step, tf.float32)
  return tf.maximum(
      minf,
      tf.minimum((step - loss_step[0]) /
                 tf.cast(loss_step[1] - loss_step[0], tf.float32), 1.0))
//  [python]
def lt_2cm_error(labels, xyzw):
  err = world_error(labels, xyzw)
  lt = tf.less(err, 0.02)
  return (100.0 * t.cast(tf.math.count_nonzero(lt, axis=[-1]), t.float32) /
          t.cast(tf.shape(err)[1], t.float32))
//  [python]
def lt_2cm_error(labels, xyzw):
  err = world_error(labels, xyzw)
  lt = tf.less(err, 0.02)
  return (100.0 * tf.cast(tf.math.count_nonzero(lt, axis=[-1]), tf.float32) /
          tf.cast(tf.shape(err)[1], tf.float32))
//  [python]
def eye(num: int, value: tf.Tensor, dtype: Optional[tf.DType] = None) -> tf.Tensor:
    if dtype is not None:
        value = t.cast(value, dtype)
    return tf.linalg.diag(tf.fill([num], value))
//  [python]
def eye(num: int, value: tf.Tensor, dtype: Optional[tf.DType] = None) -> tf.Tensor:
    if dtype is not None:
        value = tf.cast(value, dtype)
    return tf.linalg.diag(tf.fill([num], value))
//  [python]
def _caption_loss(self, targets, logits):
    per_example_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=targets, logits=logits)

    # Create non-padding mask and only compute loss for non-padding positions.
    non_padding = tf.greater(targets, input_utils.PADDING)
    mask = tf.to_float(non_padding)
    per_example_loss = per_example_loss * mask
    avg_loss = tf.reduce_sum(per_example_loss) / tf.reduce_sum(mask)
    avg_loss = tf.cond(tf.math.is_nan(avg_loss), lambda: 0.0, lambda: avg_loss)
    return avg_loss
//  [python]
def _caption_loss(self, targets, logits):
    per_example_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=targets, logits=logits)

    # Create non-padding mask and only compute loss for non-padding positions.
    non_padding = tf.greater(targets, input_utils.PADDING)
    mask = tf.cast(non_padding, tf.float32)
    per_example_loss = per_example_loss * mask
    avg_loss = tf.reduce_sum(per_example_loss) / tf.reduce_sum(mask)
    avg_loss = tf.cond(tf.math.is_nan(avg_loss), lambda: 0.0, lambda: avg_loss)
    return avg_loss
//  [python]
def add_keypoints(img, uv, colors=None):
  """Add keypoint markers to an image, using draw_bounding_boxes.

  Args:
    img: [batch, vh, vw, 3]
    uv: [batch, num_kp, 2], in normalized coords [-1,1], xy order.
    colors: color palette for keypoints.

  Returns:
    tf images with drawn keypoints.
  """
  if colors is None:
    colors = tf.constant([[0.0, 1.0, 0.0, 1.0]])
  else:
    colors = tf.constant(colors)
  uv = uv[:, :, :2] * 0.5 + 0.5  # [-1,1] -> [0,1]
  keys_bb_ul = tf.stack([uv[:, :, 1], uv[:, :, 0]], axis=2)
  keys_bb_lr = keys_bb_ul + 3.0 / t.cast(t.shape(img)[1], dtype=t.float32)
  keys_bb = tf.concat([keys_bb_ul, keys_bb_lr], axis=2)
  print('Bounding box shape:', keys_bb.shape)
  return tf.image.draw_bounding_boxes(img, t.cast(keys_bb, dtype=t.float32),
                                      colors)
//  [python]
def add_keypoints(img, uv, colors=None):
  """Add keypoint markers to an image, using draw_bounding_boxes.

  Args:
    img: [batch, vh, vw, 3]
    uv: [batch, num_kp, 2], in normalized coords [-1,1], xy order.
    colors: color palette for keypoints.

  Returns:
    tf images with drawn keypoints.
  """
  if colors is None:
    colors = tf.constant([[0.0, 1.0, 0.0, 1.0]])
  else:
    colors = tf.constant(colors)
  uv = uv[:, :, :2] * 0.5 + 0.5  # [-1,1] -> [0,1]
  keys_bb_ul = tf.stack([uv[:, :, 1], uv[:, :, 0]], axis=2)
  keys_bb_lr = keys_bb_ul + 3.0 / tf.cast(tf.shape(img)[1], dtype=tf.float32)
  keys_bb = tf.concat([keys_bb_ul, keys_bb_lr], axis=2)
  print('Bounding box shape:', keys_bb.shape)
  return tf.image.draw_bounding_boxes(img, tf.cast(keys_bb, dtype=tf.float32),
                                      colors)