code: def coordinate_tensor(shape, axis):
    
    if axis < 0:
      axis = tf.size(shape) + axis  # Convert to positive for the one_hot indice
  
    r = tf.range(shape[axis])
    r_shape = tf.one_hot(
        axis, tf.size(shape), on_value=-1, off_value=1, dtype=tf.int32)
    return tf.zeros(shape, dtype=tf.int32) + tf.reshape(r, r_shape)
docs: Return a tensor with given shape containing coordinate along given axis.
  
    Args:
      shape: a Tensor representing the shape of the output Tensor
      axis: an integer
  
    Returns:
      A tensor with shape shape and type tf.int32, where each elements its
      coordinate along the given axis.