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