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