jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def WeakSupLoss(self):
        '''
        Define additional losses for weakly supervised segmentation
        Inputs:
            self.logits = tf.reshape(self.logits, [-1, config.num_classes])
            self.labels = tf.reshape(self.labels, [-1])
            self.mask_boolen = tf.reshape(self.mask_boolen, [-1])
        Returns:

        '''
        
        ## MIL Branch
        one_hot_labels = tf.one_hot(self.batch_labels, depth=self.config.num_classes, axis=-1)
#         one_hot_labels = tf.reshape(one_hot_labels, [self.inputs.shape[0],self.inputs.shape[1],self.config.num_classes])
        L_gt = tf.cast(tf.reduce_max(one_hot_labels[0::2], axis=1), tf.float32)  # inexact labels B*13
        L_pred = tf.reduce_max(self.batch_logits[0::2], axis=1)
        loss_ineaxct = tf.nn.sigmoid_cross_entropy_with_logits(labels=L_gt, logits=L_pred)  # B*K, independent loss for each class
        self.loss_inexact = 1* tf.reduce_mean(loss_ineaxct)
        
        ## Siamese Branch
        self.prob_logits = tf.nn.softmax(self.batch_logits, axis=-1)
        self.loss_siamese = 1* tf.reduce_mean(tf.reduce_sum((self.prob_logits[0::2] - self.prob_logits[1::2]) ** 2, axis=-1))
        
#         pdb.set_trace()
        ## Smooth Branch
#         self.loss_smooth = 10* SmoothConstraint.Loss_SpatialColorSmooth_add_SelfContain(self.prob_logits[0::2], self.inputs['features'][0::2, :, 0:3], M=self.batch_masks[0::2], Y=self.batch_labels[0::2], gamma=1e-1)
        self.loss_smooth = 10 * SmoothConstraint.Loss_SpatialColorSmooth_add_SelfContain_mustLink(self.prob_logits[0::2], self.inputs['features'][0::2, :, 0:6], M=self.batch_masks[0::2], Y=self.batch_labels[0::2], gamma=1e-1)
//  [python]
def test_all_floating_input_gradient(self, dtype):
    def tf_f(x):
      res = tf.math.sin(x)
      return tf.reduce_sum(res)

    jax_f = jax2tf.call_tf(tf_f)
    tf_f_rt = jax2tf.convert(jax_f)
    x = jnp.array([5.0, 6.0, 7.0]).astype(dtype)

    def assert_all_close_support_bfloat16(baseline, candidate):
      def conversion(x):
        # convert scalar to array and bfloat16 to float32
        # to support self.assertAllClose numpy array comparison.
        if x.shape == tf.TensorShape([]):
          x = tf.convert_to_tensor([x])
        if dtype == jnp.float16:
          x = tf.cast(x, tf.float32)
        return x

      baseline = jax.tree_util.tree_map(conversion, baseline)
      candidate = jax.tree_util.tree_map(conversion, candidate)
      tol = (
          1e-2
          if jtu.test_device_matches(["tpu"]) and dtype == np.float16
          else None
      )
      self.assertAllClose(baseline, candidate, atol=tol, rtol=tol)

    # Eager mode
    assert_all_close_support_bfloat16(tf_f(x), tf_f_rt(x))

    # Compiled function mode
    assert_all_close_support_bfloat16(
        tf.function(tf_f)(x), tf.function(tf_f_rt)(x)
    )

    # Compiled function mode with jit_compiled=True
    assert_all_close_support_bfloat16(
        tf.function(tf_f, jit_compile=True)(x),
        tf.function(tf_f_rt, jit_compile=True)(x),
    )

    # RoundTrip test for the gradient
    grad_fun_jax = jax.grad(jax2tf.call_tf(tf_f))
    grad_fun_jax_rt = jax2tf.call_tf(jax2tf.convert(grad_fun_jax))

    # Eager mode
    assert_all_close_support_bfloat16(grad_fun_jax(x), grad_fun_jax_rt(x))

    # Jit mode
    assert_all_close_support_bfloat16(
        jax.jit(grad_fun_jax)(x), jax.jit(grad_fun_jax_rt)(x)
    )
//  [python]
def iou_coe(output, target, threshold=0.5, axis=[1, 2, 3], smooth=1e-5):
    """Non-differentiable Intersection over Union (IoU) for comparing the
    similarity of two batch of data, usually be used for evaluating binary image segmentation.
    The coefficient between 0 to 1, 1 means totally match.

    Parameters
    -----------
    output : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    target : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    threshold : float
        The threshold value to be true.
    axis : list of integer
        All dimensions are reduced, default ``[1,2,3]``.
    smooth : float
        This small value will be added to the numerator and denominator, see ``dice_coe``.

    Notes
    ------
    - IoU cannot be used as training loss, people usually use dice coefficient for training, IoU and hard-dice for evaluating.
    """
    pre = t.cast(output > threshold, dtype=t.float32)
    truth = t.cast(target > threshold, dtype=t.float32)
    inse = tf.reduce_sum(tf.multiply(pre, truth), axis=axis)  # AND
    union = tf.reduce_sum(tf.cast(tf.add(pre, truth) >= 1, dtype=tf.float32), axis=axis)  # OR
    ## old axis=[0,1,2,3]
    # epsilon = 1e-5
    # batch_iou = inse / (union + epsilon)
    ## new haodong
    batch_iou = (inse + smooth) / (union + smooth)
    iou = tf.reduce_mean(batch_iou)
    return iou
//  [python]
def iou_coe(output, target, threshold=0.5, axis=[1, 2, 3], smooth=1e-5):
    """Non-differentiable Intersection over Union (IoU) for comparing the
    similarity of two batch of data, usually be used for evaluating binary image segmentation.
    The coefficient between 0 to 1, 1 means totally match.

    Parameters
    -----------
    output : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    target : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    threshold : float
        The threshold value to be true.
    axis : list of integer
        All dimensions are reduced, default ``[1,2,3]``.
    smooth : float
        This small value will be added to the numerator and denominator, see ``dice_coe``.

    Notes
    ------
    - IoU cannot be used as training loss, people usually use dice coefficient for training, IoU and hard-dice for evaluating.
    """
    pre = tf.cast(output > threshold, dtype=tf.float32)
    truth = tf.cast(target > threshold, dtype=tf.float32)
    inse = tf.reduce_sum(tf.multiply(pre, truth), axis=axis)  # AND
    union = tf.reduce_sum(tf.cast(tf.add(pre, truth) >= 1, dtype=tf.float32), axis=axis)  # OR
    ## old axis=[0,1,2,3]
    # epsilon = 1e-5
    # batch_iou = inse / (union + epsilon)
    ## new haodong
    batch_iou = (inse + smooth) / (union + smooth)
    iou = tf.reduce_mean(batch_iou)
    return iou
//  [python]
def dice_hard_coe(output, target, threshold=0.5, axis=[1, 2, 3], smooth=1e-5):
    """Non-differentiable Sørensen–Dice coefficient for comparing the similarity
    of two batch of data, usually be used for binary image segmentation i.e. labels are binary.
    The coefficient between 0 to 1, 1 if totally match.

    Parameters
    -----------
    output : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    target : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    threshold : float
        The threshold value to be true.
    axis : list of integer
        All dimensions are reduced, default ``[1,2,3]``.
    smooth : float
        This small value will be added to the numerator and denominator, see ``dice_coe``.

    References
    -----------
    - `Wiki-Dice <https://en.wikipedia.org/wiki/Sørensen–Dice_coefficient>`_
    """
    output = t.cast(output > threshold, dtype=t.float32)
    target = t.cast(target > threshold, dtype=t.float32)
    inse = t.reduce_sum(t.multiply(output, target), axis=axis)
    l = t.reduce_sum(output, axis=axis)
    r = t.reduce_sum(target, axis=axis)
    ## old axis=[0,1,2,3]
    # hard_dice = 2 * (inse) / (l + r)
    # epsilon = 1e-5
    # hard_dice = tf.clip_by_value(hard_dice, 0, 1.0-epsilon)
    ## new haodong
    hard_dice = (2. * inse + smooth) / (l + r + smooth)
    ##
    hard_dice = t.reduce_mean(hard_dice)
    return hard_dice
//  [python]
def dice_hard_coe(output, target, threshold=0.5, axis=[1, 2, 3], smooth=1e-5):
    """Non-differentiable Sørensen–Dice coefficient for comparing the similarity
    of two batch of data, usually be used for binary image segmentation i.e. labels are binary.
    The coefficient between 0 to 1, 1 if totally match.

    Parameters
    -----------
    output : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    target : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    threshold : float
        The threshold value to be true.
    axis : list of integer
        All dimensions are reduced, default ``[1,2,3]``.
    smooth : float
        This small value will be added to the numerator and denominator, see ``dice_coe``.

    References
    -----------
    - `Wiki-Dice <https://en.wikipedia.org/wiki/Sørensen–Dice_coefficient>`_
    """
    output = tf.cast(output > threshold, dtype=tf.float32)
    target = tf.cast(target > threshold, dtype=tf.float32)
    inse = tf.reduce_sum(tf.multiply(output, target), axis=axis)
    l = tf.reduce_sum(output, axis=axis)
    r = tf.reduce_sum(target, axis=axis)
    ## old axis=[0,1,2,3]
    # hard_dice = 2 * (inse) / (l + r)
    # epsilon = 1e-5
    # hard_dice = tf.clip_by_value(hard_dice, 0, 1.0-epsilon)
    ## new haodong
    hard_dice = (2. * inse + smooth) / (l + r + smooth)
    ##
    hard_dice = tf.reduce_mean(hard_dice)
    return hard_dice
//  [python]
def _get_cg_discount(top_k_int=1):
  """
  Calculate discounted gain factor for ranking position till top_k_int
  Args:
    top_k_int: An int or an int `Tensor`.
  Returns:
    a `Tensor` that holds \log_{2}(i + 1), i \in [1, k] 
  """
  log_2 = tf.log(tf.constant(2.0, dtype=tf.float32))
  # top_k_range needs to start from 1 to top_k_int
  top_k_range = tf.range(top_k_int) + 1
  top_k_range = tf.reshape(top_k_range, [-1, 1])
  # cast top_k_range to float
  top_k_range = tf.to_float(top_k_range)
  cg_discount = tf.log(top_k_range + 1.0) / log_2
  return cg_discount
//  [python]
def _get_cg_discount(top_k_int=1):
  """
  Calculate discounted gain factor for ranking position till top_k_int
  Args:
    top_k_int: An int or an int `Tensor`.
  Returns:
    a `Tensor` that holds \log_{2}(i + 1), i \in [1, k] 
  """
  log_2 = tf.log(tf.constant(2.0, dtype=tf.float32))
  # top_k_range needs to start from 1 to top_k_int
  top_k_range = tf.range(top_k_int) + 1
  top_k_range = tf.reshape(top_k_range, [-1, 1])
  # cast top_k_range to float
  top_k_range = tf.cast(top_k_range, dtype=tf.float32)
  cg_discount = tf.log(top_k_range + 1.0) / log_2
  return cg_discount
//  [python]
def call(self, input_coords, input_shape=None):
        """Positionally encode points that are normalized to [0,1]."""
        coordinates = tf.identity(input_coords)

        if input_shape is not None:
            coordinates = tf.stack(
                [
                    tf.cast(coordinates[:, :, :, 0], tf.float32) / input_shape[1],
                    tf.cast(coordinates[:, :, :, 1], tf.float32) / input_shape[0],
                ],
                axis=-1,
            )

        # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
        coordinates = 2 * coordinates - 1
        coordinates = tf.cast(coordinates, self.positional_embedding.dtype)
        coordinates = tf.matmul(coordinates, self.positional_embedding)
        coordinates = 2 * np.pi * coordinates
        # outputs d_1 x ... x d_n x channel shape
        return tf.concat([tf.sin(coordinates), tf.cos(coordinates)], axis=-1)
//  [python]
def cal_err(label_scores, predicted_scores, top_k_int=1, use_predicted_order=False):
  """
  Calculate NDCG score for top_k_int ranking positions
  Args:
    label_scores: a real `Tensor`.
    predicted_scores: a real `Tensor`, with dtype matching label_scores
    top_k_int: An int or an int `Tensor`.
  Returns:
    a `Tensor` that holds DCG / IDCG.
  """
    
  if not use_predicted_order:
    max_label_score = tf.reduce_max(label_scores)
    sorted_labels, predicted_order, sorted_predictions = _get_ranking_orders(
      label_scores, predicted_scores, top_k_int=top_k_int)
  else:
    indices = tf.range(top_k_int)
    predicted_order = tf.gather(predicted_scores, indices)
    max_label_score = tf.reduce_max(label_scores)

  predicted_relevance = _get_relevance_scores(predicted_order)
  predicted_relevance_ratio = predicted_relevance / (2**max_label_score)
  _ratio = 1- predicted_relevance_ratio
  prob_stepdown = tf.cumprod(_ratio, exclusive=True)
  cur_rank = tf.range(top_k_int) + 1
  cur_rank = tf.cast(cur_rank, dtype=tf.float32)
  cur_rank = tf.reshape(cur_rank, [-1, 1])
  products = tf.multiply(predicted_relevance_ratio, prob_stepdown)
  products = tf.cast(products, dtype=tf.float32)
  errs = tf.divide(products, cur_rank)
  err = tf.reduce_sum(errs)
  return err
//  [python]
def query_stats(self, adj_matrix, cl_adj_matrix, logits, n_clauses, n_vars, query, step: str = "last"):
        current_labels = tf.round(tf.sigmoid(logits))
        round_query = tf.round(tf.sigmoid(query))
        query_logits_match = tf.equal(current_labels, round_query)

        lit_in_n_clauses = tf.sparse.reduce_sum(adj_matrix, axis=-1)
        lit1, lit2 = tf.split(lit_in_n_clauses, 2, axis=0)
        vars_count = lit1 + lit2
        in_n_clauses_matched = tf.reduce_mean(tf.expand_dims(vars_count, axis=-1) * tf.cast(query_logits_match, tf.float32))
        in_n_clauses_not_matched = tf.reduce_mean(tf.expand_dims(vars_count, axis=-1) * tf.cast(1 - query_logits_match, tf.float32))

        query_matching_values = tf.sigmoid(query) * tf.cast(1 - query_logits_match, tf.float32)
        query_not_matching_values = tf.sigmoid(query) * tf.cast(1 - query_logits_match, tf.float32)
        not_matching_mean = tf.reduce_mean(query_not_matching_values)
        not_matching_median = tfp.stats.percentile(query_not_matching_values, 50.0, interpolation='midpoint')

        matching_mean = tf.reduce_mean(query_matching_values)
        matching_median = tfp.stats.percentile(query_matching_values, 50.0, interpolation='midpoint')

        query_logits_match = tf.reduce_sum(tf.cast(query_logits_match, tf.float32), axis=0) / tf.cast(n_vars, tf.float32)
        query_logits_match = tf.reduce_mean(query_logits_match)

        lit = tf.concat([round_query, 1 - round_query], axis=0)
        sat_clauses = tf.sparse.sparse_dense_matmul(cl_adj_matrix, lit)
        sat_clauses = tf.clip_by_value(sat_clauses, 0, 1)
        sat_clauses = tf.reduce_sum(sat_clauses, axis=0) / tf.cast(n_clauses, tf.float32)
        sat_clauses = tf.reduce_mean(sat_clauses)

        with tf.name_scope(f"query_stats_{step}"):
            tf.summary.scalar("query_logits_match", query_logits_match)
            tf.summary.scalar("sat_clauses", sat_clauses)

            tf.summary.scalar("vars_in_clauses_matched", in_n_clauses_matched)
            tf.summary.scalar("vars_in_clauses_not_matched", in_n_clauses_not_matched)

            tf.summary.scalar("not_matching_mean", not_matching_mean)
            tf.summary.scalar("not_matching_median", not_matching_median)

            tf.summary.scalar("matching_mean", matching_mean)
            tf.summary.scalar("matching_median", matching_median)
//  [python]
def query_stats(self, adj_matrix, cl_adj_matrix, logits, n_clauses, n_vars, query, step: str = "last"):
        current_labels = tf.round(tf.sigmoid(logits))
        round_query = tf.round(tf.sigmoid(query))
        query_logits_match = tf.cast(tf.equal(current_labels, round_query), tf.float32)

        lit_in_n_clauses = tf.sparse.reduce_sum(adj_matrix, axis=-1)
        lit1, lit2 = tf.split(lit_in_n_clauses, 2, axis=0)
        vars_count = lit1 + lit2
        in_n_clauses_matched = tf.reduce_mean(tf.expand_dims(vars_count, axis=-1) * query_logits_match)
        in_n_clauses_not_matched = tf.reduce_mean(tf.expand_dims(vars_count, axis=-1) * (1 - query_logits_match))

        query_matching_values = tf.sigmoid(query) * (1 - query_logits_match)
        query_not_matching_values = tf.sigmoid(query) * (1 - query_logits_match)
        not_matching_mean = tf.reduce_mean(query_not_matching_values)
        not_matching_median = tfp.stats.percentile(query_not_matching_values, 50.0, interpolation='midpoint')

        matching_mean = tf.reduce_mean(query_matching_values)
        matching_median = tfp.stats.percentile(query_matching_values, 50.0, interpolation='midpoint')

        query_logits_match = tf.reduce_sum(query_logits_match, axis=0) / tf.cast(n_vars, tf.float32)
        query_logits_match = tf.reduce_mean(query_logits_match)

        lit = tf.concat([round_query, 1 - round_query], axis=0)
        sat_clauses = tf.sparse.sparse_dense_matmul(cl_adj_matrix, lit)
        sat_clauses = tf.clip_by_value(sat_clauses, 0, 1)
        sat_clauses = tf.reduce_sum(sat_clauses, axis=0) / tf.cast(n_clauses, tf.float32)
        sat_clauses = tf.reduce_mean(sat_clauses)

        with tf.name_scope(f"query_stats_{step}"):
            tf.summary.scalar("query_logits_match", query_logits_match)
            tf.summary.scalar("sat_clauses", sat_clauses)

            tf.summary.scalar("vars_in_clauses_matched", in_n_clauses_matched)
            tf.summary.scalar("vars_in_clauses_not_matched", in_n_clauses_not_matched)

            tf.summary.scalar("not_matching_mean", not_matching_mean)
            tf.summary.scalar("not_matching_median", not_matching_median)

            tf.summary.scalar("matching_mean", matching_mean)
            tf.summary.scalar("matching_median", matching_median)
//  [python]
def _to_tensor(x, dtype):
    """Convert the input `x` to a tensor of type `dtype`.

    # Arguments
        x: An object to be converted (numpy array, list, tensors).
        dtype: The destination type.

    # Returns
        A tensor.
    """
    x = tf.convert_to_tensor(x)
    if x.dtype != dtype:
        x = tf.to_float(x)
    return x
//  [python]
def _to_tensor(x, dtype):
    """Convert the input `x` to a tensor of type `dtype`.

    # Arguments
        x: An object to be converted (numpy array, list, tensors).
        dtype: The destination type.

    # Returns
        A tensor.
    """
    x = tf.convert_to_tensor(x)
    if x.dtype != dtype:
        x = tf.cast(x, dtype)
    return x
//  [python]
def _preprocess_conv2d_kernel(kernel, data_format):
    """Transpose and cast the kernel before the conv2d.

    # Arguments
        kernel: kernel tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

    # Returns
        A tensor.
    """
    if dtype(kernel) == 'float64':
        kernel = t.cast(kernel, 'float32')
    if data_format == 'channels_first':
        kernel = tf.transpose(kernel, (2, 3, 1, 0))
    return kernel
//  [python]
def _preprocess_conv2d_kernel(kernel, data_format):
    """Transpose and cast the kernel before the conv2d.

    # Arguments
        kernel: kernel tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

    # Returns
        A tensor.
    """
    if dtype(kernel) == 'float64':
        kernel = tf.cast(kernel, 'float32')
    if data_format == 'channels_first':
        kernel = tf.transpose(kernel, (2, 3, 1, 0))
    return kernel
//  [python]
def _postprocess_conv2d_output(x, data_format):
    """Transpose and cast the output from conv2d if needed.

    # Arguments
        x: A tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

    # Returns
        A tensor.
    """

    if data_format == 'channels_first':
        x = tf.transpose(x, (0, 3, 1, 2))

    if floatx() == 'float64':
        x = tf.cast(x, 'float64')
    return x
//  [python]
def switch(condition, then_expression, else_expression):
    """Switches between two operations depending on a scalar value.

    Note that both `then_expression` and `else_expression`
    should be symbolic tensors of the *same shape*.

    # Arguments
        condition: scalar tensor (`int` or `bool`).
        then_expression: either a tensor, or a callable that returns a tensor.
        else_expression: either a tensor, or a callable that returns a tensor.

    # Returns
        The selected tensor.
    """
    if condition.dtype != tf.bool:
        condition = t.cast(condition, 'bool')
    if not callable(then_expression):
        def then_expression_fn():
            return then_expression
    else:
        then_expression_fn = then_expression
    if not callable(else_expression):
        def else_expression_fn():
            return else_expression
    else:
        else_expression_fn = else_expression
    x = tf.cond(condition,
                then_expression_fn,
                else_expression_fn)
    return x
//  [python]
def switch(condition, then_expression, else_expression):
    """Switches between two operations depending on a scalar value.

    Note that both `then_expression` and `else_expression`
    should be symbolic tensors of the *same shape*.

    # Arguments
        condition: scalar tensor (`int` or `bool`).
        then_expression: either a tensor, or a callable that returns a tensor.
        else_expression: either a tensor, or a callable that returns a tensor.

    # Returns
        The selected tensor.
    """
    if condition.dtype != tf.bool:
        condition = tf.cast(condition, 'bool')
    if not callable(then_expression):
        def then_expression_fn():
            return then_expression
    else:
        then_expression_fn = then_expression
    if not callable(else_expression):
        def else_expression_fn():
            return else_expression
    else:
        else_expression_fn = else_expression
    x = tf.cond(condition,
                then_expression_fn,
                else_expression_fn)
    return x
//  [python]
def _postprocess_conv3d_output(x, data_format):
    """Transpose and cast the output from conv3d if needed.

    # Arguments
        x: A tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

    # Returns
        A tensor.
    """
    if data_format == 'channels_first':
        x = tf.transpose(x, (0, 4, 1, 2, 3))

    if floatx() == 'float64':
        x = t.cast(x, 'float64')
    return x
//  [python]
def _postprocess_conv3d_output(x, data_format):
    """Transpose and cast the output from conv3d if needed.

    # Arguments
        x: A tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

    # Returns
        A tensor.
    """
    if data_format == 'channels_first':
        x = tf.transpose(x, (0, 4, 1, 2, 3))

    if floatx() == 'float64':
        x = tf.cast(x, 'float64')
    return x
//  [python]
def _preprocess_conv3d_kernel(kernel, data_format):
    """Transpose and cast the kernel before the conv3d.

    # Arguments
        kernel: kernel tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

    # Returns
        A tensor.
    """
    if dtype(kernel) == 'float64':
        kernel = t.cast(kernel, 'float32')
    if data_format == 'channels_first':
        kernel = tf.transpose(kernel, (2, 3, 4, 1, 0))
    return kernel
//  [python]
def _preprocess_conv3d_kernel(kernel, data_format):
    """Transpose and cast the kernel before the conv3d.

    # Arguments
        kernel: kernel tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

    # Returns
        A tensor.
    """
    if dtype(kernel) == 'float64':
        kernel = tf.cast(kernel, 'float32')
    if data_format == 'channels_first':
        kernel = tf.transpose(kernel, (2, 3, 4, 1, 0))
    return kernel
//  [python]
def get_valid_anchor_weights_in_flattened_image(true_image_shapes, height,
                                                width):
  """Computes valid anchor weights for an image assuming pixels to be flattened.

  This function is useful when we only want to penalize valid areas in the
  image in the case when padding is used. The function assumes that the loss
  function will be applied after flattening the spatial dimensions and returns
  anchor weights accordingly.

  Args:
    true_image_shapes: An integer tensor of shape [batch_size, 3] representing
      the true image shape (without padding) for each sample in the batch.
    height: height of the prediction from the network.
    width: width of the prediction from the network.

  Returns:
    valid_anchor_weights: a float tensor of shape [batch_size, height * width]
    with 1s in locations where the spatial coordinates fall within the height
    and width in true_image_shapes.
  """

  indices = tf.reshape(tf.range(height * width), [1, -1])
  batch_size = tf.shape(true_image_shapes)[0]
  batch_indices = tf.ones((batch_size, 1), dtype=tf.int32) * indices

  y_coords, x_coords, _ = get_row_col_channel_indices_from_flattened_indices(
      batch_indices, width, 1)

  max_y, max_x = true_image_shapes[:, 0], true_image_shapes[:, 1]
  max_x = tf.expand_dims(max_x, 1)
  max_y = tf.expand_dims(max_y, 1)

  x_coords = tf.to_float(x_coords)
  y_coords = tf.to_float(y_coords)

  valid_mask = tf.math.logical_and(x_coords < max_x, y_coords < max_y)

  return tf.to_float(valid_mask)
//  [python]
def get_valid_anchor_weights_in_flattened_image(true_image_shapes, height,
                                                width):
  """Computes valid anchor weights for an image assuming pixels to be flattened.

  This function is useful when we only want to penalize valid areas in the
  image in the case when padding is used. The function assumes that the loss
  function will be applied after flattening the spatial dimensions and returns
  anchor weights accordingly.

  Args:
    true_image_shapes: An integer tensor of shape [batch_size, 3] representing
      the true image shape (without padding) for each sample in the batch.
    height: height of the prediction from the network.
    width: width of the prediction from the network.

  Returns:
    valid_anchor_weights: a float tensor of shape [batch_size, height * width]
    with 1s in locations where the spatial coordinates fall within the height
    and width in true_image_shapes.
  """

  indices = tf.reshape(tf.range(height * width), [1, -1])
  batch_size = tf.shape(true_image_shapes)[0]
  batch_indices = tf.ones((batch_size, 1), dtype=tf.int32) * indices

  y_coords, x_coords, _ = get_row_col_channel_indices_from_flattened_indices(
      batch_indices, width, 1)

  max_y, max_x = true_image_shapes[:, 0], true_image_shapes[:, 1]
  max_x = tf.cast(tf.expand_dims(max_x, 1), tf.float32)
  max_y = tf.cast(tf.expand_dims(max_y, 1), tf.float32)

  x_coords = tf.cast(x_coords, tf.float32)
  y_coords = tf.cast(y_coords, tf.float32)

  valid_mask = tf.math.logical_and(x_coords < max_x, y_coords < max_y)

  return tf.cast(valid_mask, tf.float32)
//  [python]
def call(self, inputs, training=None, mask=None, return_crf_internal=False):        
        base_model_outputs = self.base_model(inputs, training, mask)

        # change next line, if your model has more outputs
        crf_input = base_model_outputs

        decode_sequence, potentials, sequence_length, kernel = self.crf_layer(crf_input)

        # change next line, if your base model has more outputs
        # Aways keep `(potentials, sequence_length, kernel), decode_sequence, `
        # as first two outputs of model.
        # current `self.train_step()` expected such settings
        if not tf.executing_eagerly():
            decode_sequence = tf.to_float(decode_sequence)
        outputs = (potentials, sequence_length, kernel), decode_sequence

        '''
        tf.print(potentials)
        tf.print(sequence_length)
        tf.print(kernel)

        # See https://stats.stackexchange.com/a/517325

        # compute the normalization log Z(x) of the CRF from input, lenghts and 
        # transition matrix (chain kernel)
        normalization = crf_log_norm(potentials, sequence_length, kernel)
        #sequence_length_float32 = tf.cast(sequence_length, tf.float32)
        #normalization = tf.divide(normalization, sequence_length_float32)
        #tf.print(normalization)
        
        normalization = tf.expand_dims(normalization, axis=1)
        normalization = tf.expand_dims(normalization, axis=1)

        # then compute the difference
        probabilities = tf.math.subtract(potentials, normalization)

        # and finally exponentiate to get probabilities
        probabilities = tf.math.exp(potentials)
        #tf.print(probabilities)
        '''

        if return_crf_internal:
            return outputs
        else:
            # outputs[0] is the crf internal, skip it
            output_without_crf_internal = outputs[1:]

            # it is nicer to return a tensor instead of an one tensor list
            if len(output_without_crf_internal) == 1:
                return output_without_crf_internal[0]
            else:
                return output_without_crf_internal
//  [python]
def call(self, inputs, training=None, mask=None, return_crf_internal=False):        
        base_model_outputs = self.base_model(inputs, training, mask)

        # change next line, if your model has more outputs
        crf_input = base_model_outputs

        decode_sequence, potentials, sequence_length, kernel = self.crf_layer(crf_input)

        # change next line, if your base model has more outputs
        # Aways keep `(potentials, sequence_length, kernel), decode_sequence, `
        # as first two outputs of model.
        # current `self.train_step()` expected such settings
        if not tf.executing_eagerly():
            decode_sequence = tf.cast(decode_sequence, tf.float32)
        outputs = (potentials, sequence_length, kernel), decode_sequence

        '''
        tf.print(potentials)
        tf.print(sequence_length)
        tf.print(kernel)

        # See https://stats.stackexchange.com/a/517325

        # compute the normalization log Z(x) of the CRF from input, lenghts and 
        # transition matrix (chain kernel)
        normalization = crf_log_norm(potentials, sequence_length, kernel)
        #sequence_length_float32 = tf.cast(sequence_length, tf.float32)
        #normalization = tf.divide(normalization, sequence_length_float32)
        #tf.print(normalization)
        
        normalization = tf.expand_dims(normalization, axis=1)
        normalization = tf.expand_dims(normalization, axis=1)

        # then compute the difference
        probabilities = tf.math.subtract(potentials, normalization)

        # and finally exponentiate to get probabilities
        probabilities = tf.math.exp(potentials)
        #tf.print(probabilities)
        '''

        if return_crf_internal:
            return outputs
        else:
            # outputs[0] is the crf internal, skip it
            output_without_crf_internal = outputs[1:]

            # it is nicer to return a tensor instead of an one tensor list
            if len(output_without_crf_internal) == 1:
                return output_without_crf_internal[0]
            else:
                return output_without_crf_internal
//  [python]
def between(x, low, high, dtype=tf.float32):
    a = t.cast(tf.greater_equal(x, low), dtype=dtype)
    b = t.cast(tf.less_equal(x, high), dtype=dtype)
    return a*b
//  [python]
def between(x, low, high, dtype=tf.float32):
    a = tf.cast(tf.greater_equal(x, low), dtype=dtype)
    b = tf.cast(tf.less_equal(x, high), dtype=dtype)
    return a*b
//  [python]
def get_mask(x, value=0.0):
    """
    Creates a [batch_size x max_length] mask.
    :param value: value for the mask
    :return: [batch_size x max_length] mask of zeros and "value"s
    """
    mask = tf.greater(x, 0)
    mask = t.cast(mask, tf.float32)
    # otherwise we return a boolean mask
    if value != 0.0:
        mask *= value
    return mask
//  [python]
def get_mask(x, value=0.0):
    """
    Creates a [batch_size x max_length] mask.
    :param value: value for the mask
    :return: [batch_size x max_length] mask of zeros and "value"s
    """
    mask = tf.greater(x, 0)
    mask = tf.cast(mask, tf.float32)
    # otherwise we return a boolean mask
    if value != 0.0:
        mask *= value
    return mask
//  [python]
def get_median2(sorted_numbers, sorted_weights, alpha=0.5, dtype=tf.float32):
    is_num = tf.cast(tf.greater(sorted_weights, 0.0), dtype=dtype)
    upward = tf.cumsum(sorted_weights, axis=-1, exclusive=True)
    downward = tf.cumsum(sorted_weights, axis=-1, exclusive=True, reverse=True)
    c = tf.cumsum(is_num, axis=-1, exclusive=True, reverse=True)
    loss = (upward+downward-c*sorted_numbers)
    loss = apply_mask(loss, is_num, -np.inf)
    loss = tf.nn.softmax(-loss * sorted_weights)
    median = tf.reduce_sum(loss * sorted_numbers, axis=-1)
    return median
//  [python]
def get_median(sorted_numbers, sorted_weights, alpha=0.5, dtype=tf.float32):
    total_weight = tf.reduce_sum(sorted_weights, axis=-1, keep_dims=True)
    upper_mask = tf.greater_equal(tf.cumsum(sorted_weights, axis=-1), total_weight*alpha)
    lower_mask = tf.greater_equal(tf.cumsum(sorted_weights, axis=-1, reverse=True), total_weight*alpha)
    median_mask = t.cast(upper_mask, dtype=dtype) * t.cast(lower_mask, dtype=dtype)
    median_mask *= sorted_weights  # FIXME ???
    median_mask = median_mask/tf.maximum(tf.reduce_sum(median_mask, axis=-1, keep_dims=True), 0.0000000001)
    median = tf.reduce_sum(median_mask * sorted_numbers, axis=-1)
    return median
//  [python]
def get_median(sorted_numbers, sorted_weights, alpha=0.5, dtype=tf.float32):
    total_weight = tf.reduce_sum(sorted_weights, axis=-1, keep_dims=True)
    upper_mask = tf.greater_equal(tf.cumsum(sorted_weights, axis=-1), total_weight*alpha)
    lower_mask = tf.greater_equal(tf.cumsum(sorted_weights, axis=-1, reverse=True), total_weight*alpha)
    median_mask = tf.cast(upper_mask, dtype=dtype) * tf.cast(lower_mask, dtype=dtype)
    median_mask *= sorted_weights  # FIXME ???
    median_mask = median_mask/tf.maximum(tf.reduce_sum(median_mask, axis=-1, keep_dims=True), 0.0000000001)
    median = tf.reduce_sum(median_mask * sorted_numbers, axis=-1)
    return median
//  [python]
def get_robust_statistics(x):
    # median and median absolute deviation (MAD)
    def _get_median(_x):
        # FIXME this is just an approximation
        n_elements = t.unstack(t.shape(x))[0]
        n_halfway = t.cast(t.floor(t.cast(n_elements, dtype=t.float32)/2.0), dtype=t.int32)
        topk_values, _ = t.nn.top_k(_x, k=n_halfway, sorted=True)
        _median = t.slice(topk_values, [n_halfway-1], [1])
        return _median
    median = _get_median(x)
    mad = _get_median(t.abs(x-median))
    return median, mad
//  [python]
def get_robust_statistics(x):
    # median and median absolute deviation (MAD)
    def _get_median(_x):
        # FIXME this is just an approximation
        n_elements = tf.unstack(tf.shape(x))[0]
        n_halfway = tf.cast(tf.floor(tf.cast(n_elements, dtype=tf.float32)/2.0), dtype=tf.int32)
        topk_values, _ = tf.nn.top_k(_x, k=n_halfway, sorted=True)
        _median = tf.slice(topk_values, [n_halfway-1], [1])
        return _median
    median = _get_median(x)
    mad = _get_median(tf.abs(x-median))
    return median, mad
//  [python]
def make_hard(probs, is_binary=False, dtype=tf.float32):
    if is_binary:
        probs = tf.greater_equal(probs, 0.5)
        probs = t.cast(probs, dtype=dtype)
    else:
        probs = t.cast(tf.equal(probs, tf.reduce_max(probs, axis=-1, keep_dims=True)), dtype=dtype)
        probs /= tf.reduce_sum(probs, axis=-1, keep_dims=True)  # in case of ties
    return probs
//  [python]
def make_hard(probs, is_binary=False, dtype=tf.float32):
    if is_binary:
        probs = tf.greater_equal(probs, 0.5)
        probs = tf.cast(probs, dtype=dtype)
    else:
        probs = tf.cast(tf.equal(probs, tf.reduce_max(probs, axis=-1, keep_dims=True)), dtype=dtype)
        probs /= tf.reduce_sum(probs, axis=-1, keep_dims=True)  # in case of ties
    return probs
//  [python]
def get_cumulative_attributes(x,  # (batch, time)
                              mu,  # (vocab,)
                              dtype=tf.float32
                              ):
    x = tf.expand_dims(x, -1)  # (batch, time, 1)
    mu = tf.expand_dims(tf.expand_dims(mu, 0), 0)  # (1, 1, vocab)
    attributes = tf.less_equal(mu, x)
    return tf.to_float(attributes)
//  [python]
def get_cumulative_attributes(x,  # (batch, time)
                              mu,  # (vocab,)
                              dtype=tf.float32
                              ):
    x = tf.expand_dims(x, -1)  # (batch, time, 1)
    mu = tf.expand_dims(tf.expand_dims(mu, 0), 0)  # (1, 1, vocab)
    attributes = tf.less_equal(mu, x)
    return tf.cast(attributes, dtype=dtype)
//  [python]
def build_graph(self, feat, labelidx, labelvalue, labelshape, seqlen):
        label = tf.SparseTensor(labelidx, labelvalue, labelshape)

        cell = rnn.MultiRNNCell([rnn.LSTMBlockCell(num_units=HIDDEN) for _ in range(NLAYER)])

        initial = cell.zero_state(tf.shape(feat)[0], tf.float32)

        outputs, last_state = tf.nn.dynamic_rnn(cell, feat,
                                                seqlen, initial,
                                                dtype=tf.float32, scope='rnn')

        # o: b x t x HIDDEN
        output = tf.reshape(outputs, [-1, HIDDEN])  # (Bxt) x rnnsize
        logits = FullyConnected('fc', output, NR_CLASS, activation=tf.identity,
                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.01))
        logits = tf.reshape(logits, (BATCH, -1, NR_CLASS))

        loss = tf.nn.ctc_loss(label, logits, seqlen, time_major=False)

        cost = tf.reduce_mean(loss, name='cost')

        logits = tf.transpose(logits, [1, 0, 2])

        if self.training:
            # beam search is too slow to run in training
            predictions = tf.to_int32(
                tf.nn.ctc_greedy_decoder(logits, seqlen)[0][0])
        else:
            predictions = tf.to_int32(
                tf.nn.ctc_beam_search_decoder(logits, seqlen)[0][0])
        err = tf.edit_distance(predictions, label, normalize=True)
        err.set_shape([None])
        err = tf.reduce_mean(err, name='error')
        summary.add_moving_summary(err, cost)
        return cost
//  [python]
def build_graph(self, feat, labelidx, labelvalue, labelshape, seqlen):
        label = tf.SparseTensor(labelidx, labelvalue, labelshape)

        cell = rnn.MultiRNNCell([rnn.LSTMBlockCell(num_units=HIDDEN) for _ in range(NLAYER)])

        initial = cell.zero_state(tf.shape(feat)[0], tf.float32)

        outputs, last_state = tf.nn.dynamic_rnn(cell, feat,
                                                seqlen, initial,
                                                dtype=tf.float32, scope='rnn')

        # o: b x t x HIDDEN
        output = tf.reshape(outputs, [-1, HIDDEN])  # (Bxt) x rnnsize
        logits = FullyConnected('fc', output, NR_CLASS, activation=tf.identity,
                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.01))
        logits = tf.reshape(logits, (BATCH, -1, NR_CLASS))

        loss = tf.nn.ctc_loss(label, logits, seqlen, time_major=False)

        cost = tf.reduce_mean(loss, name='cost')

        logits = tf.transpose(logits, [1, 0, 2])

        if self.training:
            # beam search is too slow to run in training
            predictions = tf.cast(
                tf.nn.ctc_greedy_decoder(logits, seqlen)[0][0], tf.int32)
        else:
            predictions = tf.cast(
                tf.nn.ctc_beam_search_decoder(logits, seqlen)[0][0], tf.int32)
        err = tf.edit_distance(predictions, label, normalize=True)
        err.set_shape([None])
        err = tf.reduce_mean(err, name='error')
        summary.add_moving_summary(err, cost)
        return cost
//  [python]
def get_cumulative_attributes2(x,  # (batch, time)
                              mu,  # (vocab,)
                              dtype=tf.float32
                              ):
    x = tf.expand_dims(x, -1)  # (batch, time, 1)
    mu = tf.expand_dims(tf.expand_dims(mu, 0), 0)  # (1, 1, vocab)
    dist = tf.abs(mu-x)
    attributes = tf.equal(dist, tf.reduce_min(dist,axis=-1,keep_dims=True))
    return tf.to_float(attributes)
//  [python]
def get_cumulative_attributes2(x,  # (batch, time)
                              mu,  # (vocab,)
                              dtype=tf.float32
                              ):
    x = tf.expand_dims(x, -1)  # (batch, time, 1)
    mu = tf.expand_dims(tf.expand_dims(mu, 0), 0)  # (1, 1, vocab)
    dist = tf.abs(mu-x)
    attributes = tf.equal(dist, tf.reduce_min(dist,axis=-1,keep_dims=True))
    return tf.cast(attributes, dtype=dtype)
//  [python]
def mi_to_tf_tensor(mi_tensor, dtype):
    """
    Get a TensorFlow eager tensor from a Mitsuba/DrJIT tensor
    """
    # When there is only one input, the .tf() methods crashes.
    # The following hack takes care of this corner case
    dr.eval(mi_tensor)
    dr.sync_thread()
    if dr.shape(mi_tensor)[-1] == 1:
        mi_tensor = dr.repeat(mi_tensor, 2)
        tf_tensor = t.cast(mi_tensor.tf(), dtype)[:1]
    else:
        tf_tensor = t.cast(mi_tensor.tf(), dtype)
    return tf_tensor
//  [python]
def mi_to_tf_tensor(mi_tensor, dtype):
    """
    Get a TensorFlow eager tensor from a Mitsuba/DrJIT tensor
    """
    # When there is only one input, the .tf() methods crashes.
    # The following hack takes care of this corner case
    dr.eval(mi_tensor)
    dr.sync_thread()
    if dr.shape(mi_tensor)[-1] == 1:
        mi_tensor = dr.repeat(mi_tensor, 2)
        tf_tensor = tf.cast(mi_tensor.tf(), dtype)[:1]
    else:
        tf_tensor = tf.cast(mi_tensor.tf(), dtype)
    return tf_tensor
//  [python]
def sign(x):
    """
    Returns +1 if ``x`` is non-negative, -1 otherwise

    Input
    ------
    x : [...], tf.float
        A real-valued number

    Output
    -------
    : [...], tf.float
        +1 if ``x`` is non-negative, -1 otherwise
    """
    two = t.cast(2, x.dtype)
    one = t.cast(1, x.dtype)
    return two*t.cast(tf.greater_equal(x, 0), x.dtype)-one
//  [python]
def sign(x):
    """
    Returns +1 if ``x`` is non-negative, -1 otherwise

    Input
    ------
    x : [...], tf.float
        A real-valued number

    Output
    -------
    : [...], tf.float
        +1 if ``x`` is non-negative, -1 otherwise
    """
    two = tf.cast(2, x.dtype)
    one = tf.cast(1, x.dtype)
    return two*tf.cast(tf.greater_equal(x, 0), x.dtype)-one
//  [python]
def gen_orthogonal_vector(k, epsilon):
    """
    Generate an arbitrary vector that is orthogonal to ``k``.

    Input
    ------
    k : [..., 3], tf.float
        Vector

    epsilon : (), tf.float
        Small value used to avoid errors due to numerical precision

    Output
    -------
    : [..., 3], tf.float
        Vector orthogonal to ``k``
    """
    rdtype = k.dtype
    ex = t.cast([1.0, 0.0, 0.0], rdtype)
    ex = expand_to_rank(ex, tf.rank(k), 0)

    ey = t.cast([0.0, 1.0, 0.0], rdtype)
    ey = expand_to_rank(ey, tf.rank(k), 0)

    n1 = cross(k, ex)
    n1_norm = tf.norm(n1, axis=-1, keepdims=True)
    n2 = cross(k, ey)
    return tf.where(tf.greater(n1_norm, epsilon), n1, n2)
//  [python]
def gen_orthogonal_vector(k, epsilon):
    """
    Generate an arbitrary vector that is orthogonal to ``k``.

    Input
    ------
    k : [..., 3], tf.float
        Vector

    epsilon : (), tf.float
        Small value used to avoid errors due to numerical precision

    Output
    -------
    : [..., 3], tf.float
        Vector orthogonal to ``k``
    """
    rdtype = k.dtype
    ex = tf.cast([1.0, 0.0, 0.0], rdtype)
    ex = expand_to_rank(ex, tf.rank(k), 0)

    ey = tf.cast([0.0, 1.0, 0.0], rdtype)
    ey = expand_to_rank(ey, tf.rank(k), 0)

    n1 = cross(k, ex)
    n1_norm = tf.norm(n1, axis=-1, keepdims=True)
    n2 = cross(k, ey)
    return tf.where(tf.greater(n1_norm, epsilon), n1, n2)
//  [python]
def fibonacci_lattice(num_points, dtype=tf.float32):
    """
    Generates a Fibonacci lattice for the unit square

    Input
    -----
    num_points : int
        Number of points

    type : tf.DType
        Datatype to use for the output

    Output
    -------
    points : [num_points, 2]
        Generated rectangular coordinates of the lattice points
    """

    golden_ratio = (1.+tf.sqrt(tf.cast(5., tf.float64)))/2.
    ns = tf.range(0, num_points, dtype=tf.float64)

    x = ns/golden_ratio
    x = x - tf.floor(x)
    y = ns/(num_points-1)
    points = tf.stack([x,y], axis=1)

    points = tf.to_float(points)

    return points
//  [python]
def fibonacci_lattice(num_points, dtype=tf.float32):
    """
    Generates a Fibonacci lattice for the unit square

    Input
    -----
    num_points : int
        Number of points

    type : tf.DType
        Datatype to use for the output

    Output
    -------
    points : [num_points, 2]
        Generated rectangular coordinates of the lattice points
    """

    golden_ratio = (1.+tf.sqrt(tf.cast(5., tf.float64)))/2.
    ns = tf.range(0, num_points, dtype=tf.float64)

    x = ns/golden_ratio
    x = x - tf.floor(x)
    y = ns/(num_points-1)
    points = tf.stack([x,y], axis=1)

    points = tf.cast(points, dtype)

    return points
//  [python]
def theta_phi_from_unit_vec(v):
    r"""
    Computes zenith and azimuth angles (:math:`\theta,\varphi`)
    from unit-norm vectors as described in :eq:`theta_phi`

    Input
    ------
    v : [...,3], tf.float
        Tensor with unit-norm vectors in the last dimension

    Output
    -------
    theta : [...], tf.float
        Zenith angles :math:`\theta`

    phi : [...], tf.float
        Azimuth angles :math:`\varphi`
    """
    x = v[...,0]
    y = v[...,1]
    z = v[...,2]

    # If v = z, then x = 0 and y = 0. In this case, atan2 is not differentiable,
    # leading to NaN when computing the gradients.
    # The following lines force x to one this case. Note that this does not
    # impact the output meaningfully, as in that case theta = 0 and phi can
    # take any value.
    zero = tf.zeros_like(x)
    is_unit_z = tf.logical_and(tf.equal(x, zero), tf.equal(y, zero))
    is_unit_z = tf.to_float(is_unit_z)
    x += is_unit_z

    theta = acos_diff(z)
    phi = tf.math.atan2(y, x)
    return theta, phi
//  [python]
def theta_phi_from_unit_vec(v):
    r"""
    Computes zenith and azimuth angles (:math:`\theta,\varphi`)
    from unit-norm vectors as described in :eq:`theta_phi`

    Input
    ------
    v : [...,3], tf.float
        Tensor with unit-norm vectors in the last dimension

    Output
    -------
    theta : [...], tf.float
        Zenith angles :math:`\theta`

    phi : [...], tf.float
        Azimuth angles :math:`\varphi`
    """
    x = v[...,0]
    y = v[...,1]
    z = v[...,2]

    # If v = z, then x = 0 and y = 0. In this case, atan2 is not differentiable,
    # leading to NaN when computing the gradients.
    # The following lines force x to one this case. Note that this does not
    # impact the output meaningfully, as in that case theta = 0 and phi can
    # take any value.
    zero = tf.zeros_like(x)
    is_unit_z = tf.logical_and(tf.equal(x, zero), tf.equal(y, zero))
    is_unit_z = tf.cast(is_unit_z, x.dtype)
    x += is_unit_z

    theta = acos_diff(z)
    phi = tf.math.atan2(y, x)
    return theta, phi
//  [python]
def length_normalization(length, dtype):
  """Return length normalization factor."""
  return tf.pow(((5. + tf.cast(length, dtype)) / 6.), 0.0)
//  [python]
def box_list_scale(boxlist, y_scale, x_scale, scope=None):
  """scale box coordinates in x and y dimensions.

  Args:
    boxlist: BoxList holding N boxes
    y_scale: (float) scalar tensor
    x_scale: (float) scalar tensor
    scope: name scope.

  Returns:
    boxlist: BoxList holding N boxes
  """
  if not scope:
    scope = 'Scale'
  with tf.name_scope(scope):
    y_scale = tf.cast(y_scale, tf.float32)
    x_scale = tf.cast(x_scale, tf.float32)
    y_min, x_min, y_max, x_max = tf.split(
        value=boxlist.get(), num_or_size_splits=4, axis=1)
    y_min = y_scale * y_min
    y_max = y_scale * y_max
    x_min = x_scale * x_min
    x_max = x_scale * x_max
    scaled_boxlist = box_list.BoxList(
        tf.concat([y_min, x_min, y_max, x_max], 1))
    return _copy_extra_fields(scaled_boxlist, boxlist)
//  [python]
def test_get_shape_keep_last_dim(self):
    y = tf.constant(4.0)
    x = tf.ones([7, tf.cast(tf.sqrt(y), tf.int32), 2, 5])
    shape = decoding_module.get_shape_keep_last_dim(x)
    self.assertAllEqual([None, None, None, 5], shape.as_list())
//  [python]
def keypoint_scale(keypoints, y_scale, x_scale, scope=None):
  """Scales keypoint coordinates in x and y dimensions.

  Args:
    keypoints: a tensor of shape [num_instances, num_keypoints, 2]
    y_scale: (float) scalar tensor
    x_scale: (float) scalar tensor
    scope: name scope.

  Returns:
    new_keypoints: a tensor of shape [num_instances, num_keypoints, 2]
  """
  if not scope:
    scope = 'Scale'
  with tf.name_scope(scope):
    y_scale = tf.to_float(y_scale)
    x_scale = tf.to_float(x_scale)
    new_keypoints = keypoints * [[[y_scale, x_scale]]]
    return new_keypoints