code: def get_batch_coordinate(x, axis=0):
    
    # Compute the batch coordinate before flattening all batches
    batch_coordinate = tf.expand_dims(
        common_attention.coordinate_tensor(tf.shape(x)[:-1], axis=axis), axis=-1)
    return batch_coordinate
docs: Return a flat int32 tensor of shape [1, batch_size*length, 1].