code: def linear_interpolate_rank(tensor1, tensor2, coeffs, rank=1):
    
    # sum across space, max across channels.
    _, _, _, num_channels = common_layers.shape_list(tensor1)
    diff_sq_sum = tf.reduce_sum((tensor1 - tensor2)**2, axis=(0, 1, 2))
    _, feature_ranks = tf.math.top_k(diff_sq_sum, k=rank)
    feature_rank = feature_ranks[-1]
    channel_inds = tf.range(num_channels, dtype=tf.int32)
    channel_mask = tf.equal(channel_inds, feature_rank)
    ones_t = tf.ones(num_channels, dtype=tf.float32)
    zeros_t = tf.zeros(num_channels, dtype=tf.float32)
  
    interp_tensors = []
    for coeff in coeffs:
      curr_coeff = tf.where(channel_mask, coeff * ones_t, zeros_t)
      interp_tensor = tensor1 + curr_coeff * (tensor2 - tensor1)
      interp_tensors.append(interp_tensor)
    return tf.concat(interp_tensors, axis=0)
docs: Linearly interpolate channel at "rank" between two tensors.
  
    The channels are ranked according to their L2 norm between tensor1[channel]
    and tensor2[channel].
  
    Args:
      tensor1: 4-D Tensor, NHWC
      tensor2: 4-D Tensor, NHWC
      coeffs: list of floats.
      rank: integer.
    Returns:
      interp_latents: list of interpolated 4-D Tensors, shape=(NHWC)