def diff(x, axis=-1):
"""Take the finite difference of a tensor along an axis.
Args:
x: Input tensor of any dimension.
axis: Axis on which to take the finite difference.
Returns:
d: Tensor with size less than x by 1 along the difference dimension.
Raises:
ValueError: Axis out of range for tensor.
"""
shape = x.shape.as_list()
if axis >= len(shape):
raise ValueError(
'Invalid axis index: %d for tensor with only %d axes.' %
(axis, len(shape)))
begin_back = [0 for _ in range(len(shape))]
begin_front = [0 for _ in range(len(shape))]
begin_front[axis] = 1
size = shape.as_list()
size[axis] -= 1
slice_front = tf.slice(x, begin_front, size)
slice_back = tf.slice(x, begin_back, size)
d = slice_front - slice_back
return d
def _parse_record(self, record) -> NestedMap:
"""Reads and parses a single record."""
p = self.params
name_to_features = {
'input_ids':
tf.io.FixedLenFeature([p.max_sequence_length], tf.int64),
'input_mask':
tf.io.FixedLenFeature([p.max_sequence_length], tf.int64),
'masked_lm_positions':
tf.io.FixedLenFeature([p.max_predictions_per_seq], tf.int64),
'masked_lm_ids':
tf.io.FixedLenFeature([p.max_predictions_per_seq], tf.int64),
'masked_lm_weights':
tf.io.FixedLenFeature([p.max_predictions_per_seq], tf.float32),
}
example = tf.io.parse_single_example(record, name_to_features)
mask_length = tf.cast(
tf.reduce_sum(example['masked_lm_weights']), dtype=tf.int32)
masked_lm_positions = tf.slice(example['masked_lm_positions'], [0],
[mask_length])
masked_lm_ids = tf.cast(
tf.slice(example['masked_lm_ids'], [0], [mask_length]), dtype=tf.int32)
ret = py_utils.NestedMap()
ret.masked_ids = tf.cast(example['input_ids'], dtype=tf.int32)
# Get back non-masked, original ids.
ret.labels = tf.tensor_scatter_nd_update(
tensor=ret.masked_ids,
indices=tf.reshape(masked_lm_positions, [-1, 1]),
updates=masked_lm_ids)
ret.masked_pos = tf.tensor_scatter_nd_update(
tensor=tf.zeros_like(ret.masked_ids, dtype=tf.float32),
indices=tf.reshape(masked_lm_positions, [-1, 1]),
updates=tf.ones_like(masked_lm_ids, dtype=tf.float32))
ret.segment_ids = tf.cast(example['input_mask'], dtype=tf.float32)
first_eos_idx = tf.where(tf.math.equal(ret.labels, p.eos_token_id))[0][0]
def remove_first_eos(x):
# We remove the element at position `first_eos_idx`, and pad with 0
# to keep length unchanged.
zero = tf.constant(0, shape=(1,), dtype=x.dtype)
return tf.concat([x[:first_eos_idx], x[first_eos_idx + 1:], zero], axis=0)
ret = ret.Transform(remove_first_eos)
ret.paddings = 1.0 - ret.segment_ids
pos = tf.cast(tf.range(p.max_sequence_length), dtype=tf.float32)
ret.segment_pos = tf.cast(ret.segment_ids * pos, dtype=tf.int32)
if p.remask:
new_masked_ids, new_masked_pos = self.mlm.FProp(None, ret.labels,
ret.paddings)
ret.masked_ids = new_masked_ids
ret.masked_pos = new_masked_pos
return ret
def eye(N, M=None, k=0, dtype=float): # pylint: disable=invalid-name
"""Returns a 2-D array with ones on the diagonal and zeros elsewhere.
Examples:
```python
eye(2, dtype=int)
-> [[1, 0],
[0, 1]]
eye(2, M=3, dtype=int)
-> [[1, 0, 0],
[0, 1, 0]]
eye(2, M=3, k=1, dtype=int)
-> [[0, 1, 0],
[0, 0, 1]]
eye(3, M=2, k=-1, dtype=int)
-> [[0, 0],
[1, 0],
[0, 1]]
```
Args:
N: integer. Number of rows in output array.
M: Optional integer. Number of cols in output array, defaults to N.
k: Optional integer. Position of the diagonal. The default 0 refers to the
main diagonal. A positive/negative value shifts the diagonal by the
corresponding positions to the right/left.
dtype: Optional, defaults to float. The type of the resulting ndarray.
Could be a python type, a NumPy type or a TensorFlow `DType`.
Returns:
An ndarray with shape (N, M) and requested type.
"""
if dtype:
dtype = utils.to_tf_type(dtype)
if not M:
M = N
if k >= M or -k >= N:
return zeros([N, M], dtype=dtype)
if k:
if k > 0:
result = tf.eye(N, M, dtype=dtype)
zero_cols = tf.zeros([N, abs(k)], dtype=dtype)
result = tf.concat([zero_cols, result], axis=1)
result = tf.slice(result, [0, 0], [N, M])
else:
result = tf.eye(N, M - k, dtype=dtype)
result = tf.slice(result, [0, -k], [N, M])
else:
result = tf.eye(N, M, dtype=dtype)
return utils.tensor_to_ndarray(result)
def _shift(tensor, axis, delta):
"""Shifts the given tensor, filling it with zeros on the other side.
Args:
tensor: `Tensor`.
axis: Axis to shift along.
delta: Shift size. May be negative: the sign determines the direction of the
shift.
Returns:
Shifted `Tensor`.
Example:
```
t = [[1, 2, 3]
[4, 5, 6]
[7, 8, 9]]
_shift(t, 1, 2) = [[0, 0, 1]
[0, 0, 4]
[0, 0, 7]]
_shift(t, 0, -1) = [[4, 5, 6]
[7, 8, 9]
[0, 0, 0]]
TODO(b/144087751): implement this in C++. Perhaps we can add a parameter to
tf.roll, so that it fills "the other side" with zeros.
"""
rank = len(tensor.shape)
zeros_shape = np.zeros(rank)
for d in range(rank):
if d == axis:
zeros_shape[d] = np.abs(delta)
else:
zeros_shape[d] = tf.compat.dimension_value(tensor.shape[d])
zeros = tf.zeros(zeros_shape, dtype=tensor.dtype)
slice_begin = np.zeros(rank, dtype=np.int32)
slice_size = -np.ones(rank, dtype=np.int32)
if delta > 0:
slice_size[axis] = tf.compat.dimension_value(
tensor.shape[axis]) - delta
return tf.concat((zeros, tf.slice(tensor, slice_begin, slice_size)),
axis=axis)
else:
slice_begin[axis] = -delta
return tf.concat((tf.slice(tensor, slice_begin, slice_size), zeros),
axis=axis)
def subsample_indicator(indicator, num_samples):
"""Subsample indicator vector.
Given a boolean indicator vector with M elements set to `True`, the function
assigns all but `num_samples` of these previously `True` elements to
`False`. If `num_samples` is greater than M, the original indicator vector
is returned.
Args:
indicator: a 1-dimensional boolean tensor indicating which elements
are allowed to be sampled and which are not.
num_samples: int32 scalar tensor
Returns:
a boolean tensor with the same shape as input (indicator) tensor
"""
indices = tf.where(indicator)
indices = tf.random.shuffle(indices)
indices = tf.reshape(indices, [-1])
num_samples = tf.minimum(tf.size(input=indices), num_samples)
selected_indices = tf.slice(indices, [0], tf.reshape(num_samples, [1]))
selected_indicator = ops.indices_to_dense_vector(
selected_indices,
tf.shape(input=indicator)[0])
return tf.equal(selected_indicator, 1)
def slice_of(xn):
slice_size = tf.where(
tf.equal(tf.range(tf.rank(xn)),
tf.rank(xn) + axis), ax_minsize, tf.shape(xn))
return tf.slice(xn,
begin=tf.zeros_like(tf.shape(xn)),
size=slice_size)
def circular_pad(input_tensor, axis, padding):
"""Pads tensor circularly.
More specifically, pad on the right with the tensor values from the left of
the tensor, as if you had concatenated the tensor on the right and vice versa.
Args:
input_tensor: typically a batch of input "images"
axis: on which to perform the circluar padding
padding: See tf.nn.conv2d arg: padding
Returns:
Tensor of shape BxHxWxC2
"""
assert 0 <= axis < len(input_tensor.shape), 'Axis out of bounds'
multiples = [1] * len(input_tensor.shape)
multiples[axis] = 3
tiled_input = tf.tile(input_tensor, multiples)
left = input_tensor.shape[axis] - padding[0]
right = 2 * input_tensor.shape[axis] + padding[1]
begin = [0] * len(input_tensor.shape)
end = list(input_tensor.shape)
begin[axis] = left
end[axis] = right
size = [a - b for a, b in zip(end, begin)]
output_tensor = tf.slice(tiled_input, begin, size)
# Do a shape assert
return output_tensor
def input_fn():
# text input
text = tf.compat.v1.placeholder(tf.string, [batch_size], name="input_text")
# text tokenize
tokenizer = tft.SentencepieceTokenizer(
model=tf.io.gfile.GFile(vocab_model_file, "rb").read())
if substitute_newline:
text = tf.strings.regex_replace(text, "\n", substitute_newline)
ids = tokenizer.tokenize(text)
ids = ids[:, :max_encoder_length - 2]
# Add [CLS] and [SEP] special tokens.
prefix = tf.repeat(tf.constant([[65]]), batch_size, axis=0)
suffix = tf.repeat(tf.constant([[66]]), batch_size, axis=0)
ids = tf.concat([prefix, ids, suffix], axis=1)
if isinstance(ids, tf.RaggedTensor):
ids = ids.to_tensor(0)
# text padding: Pad only if necessary and reshape properly
padded_ids = dynamic_padding(ids, max_encoder_length)
ids = tf.slice(padded_ids, [0, 0], [batch_size, max_encoder_length])
receiver_tensors = {"input": text}
features = {"input_ids": tf.cast(ids, tf.int32, name="input_ids")}
return tf.estimator.export.ServingInputReceiver(
features=features, receiver_tensors=receiver_tensors)
def _inverse(self, y):
ndims = prefer_static.rank(y)
indices = prefer_static.reshape(prefer_static.add(self.axis, ndims),
shape=[-1, 1])
num_left, num_right = prefer_static.unstack(self.paddings,
num=2,
axis=-1)
x = tf.slice(y,
begin=prefer_static.tensor_scatter_nd_update(
prefer_static.zeros(ndims, dtype=tf.int32), indices,
num_left),
size=prefer_static.tensor_scatter_nd_sub(
prefer_static.shape(y), indices,
num_left + num_right))
if not self.validate_args:
return x
assertions = [
assert_util.assert_equal(
self._forward(x),
y,
message=('Argument `y` to `inverse` was not padded with '
'`constant_values`.')),
]
with tf.control_dependencies(assertions):
return tf.identity(x)
def metric_fn(logits, dup_mask, match_mlperf):
dup_mask = tf.cast(dup_mask, tf.float32)
logits = tf.slice(logits, [0, 1], [-1, -1])
in_top_k, _, metric_weights, _ = neumf_model.compute_top_k_and_ndcg(
logits, dup_mask, match_mlperf)
metric_weights = tf.cast(metric_weights, tf.float32)
return in_top_k, metric_weights
def input_fn():
# text input
text = tf.compat.v1.placeholder(tf.string, [batch_size],
name="input_text")
# text tokenize
tokenizer = tft.SentencepieceTokenizer(
model=tf.io.gfile.GFile(vocab_model_file, "rb").read())
if substitute_newline:
text = tf.strings.regex_replace(text, "\n", substitute_newline)
# Remove space before special tokens.
text = tf.strings.regex_replace(text, r" ([<\[]\S+[>\]])", b"\\1")
ids = tokenizer.tokenize(text)
if isinstance(ids, tf.RaggedTensor):
ids = ids.to_tensor(0)
# text padding: Pad only if necessary and reshape properly
padded_ids = dynamic_padding(ids, max_encoder_length)
ids = tf.slice(padded_ids, [0, 0], [batch_size, max_encoder_length])
receiver_tensors = {"input": text}
features = {"input_ids": tf.cast(ids, tf.int32, name="input_ids")}
return tf.estimator.export.ServingInputReceiver(
features=features, receiver_tensors=receiver_tensors)
def _first_element(x):
if x.shape.ndims is None:
raise ValueError('Rank of Tensor %s must be known' % x)
ndims = x.shape.ndims
begin = tf.zeros(ndims, dtype=tf.int32)
size = tf.ones(ndims, dtype=tf.int32)
return tf.reshape(tf.slice(x, begin, size), [])
def basis(sample_paths, time_index):
"""Computes polynomial basis expansion at the given sample points.
Args:
sample_paths: A `Tensor` of either `flaot32` or `float64` dtype and of
shape `[num_samples, num_times, dim]`.
time_index: An integer scalar `Tensor` that corresponds to the time
coordinate at which the basis function is computed.
Returns:
A `Tensor`s of shape `[(degree + 1)**dim, num_samples]`.
"""
sample_paths = tf.convert_to_tensor(sample_paths, name="sample_paths")
shape = sample_paths.shape.as_list()
num_samples = shape[0]
dim = shape[-1]
slice_samples = tf.slice(sample_paths, [0, time_index, 0],
[num_samples, 1, dim])
slice_samples = tf.squeeze(slice_samples, 1)
dim = sample_paths.shape.as_list()[-1]
samples_centered = slice_samples - tf.math.reduce_mean(slice_samples,
axis=0)
samples_centered = tf.expand_dims(samples_centered, axis=-2)
grid = tf.range(degree + 1, dtype=samples_centered.dtype)
# Creates a grid of 'power' expansions, i.e., a `Tensor` of shape
# [(degree + 1)**dim, dim] with entries [k_1, .., k_dim] where
## 0 <= k_i <= dim.
grid = tf.meshgrid(*(dim * [grid]))
# Shape [(degree + 1)**3, dim]
grid = tf.reshape(tf.stack(grid, -1), [-1, dim])
# `samples_centered` has shape [num_samples, 1, dim],
# `samples_centered**grid` has shape `[num_samples, (degree + 1)**dim, dim]`
# so that the output shape is [num_samples, (degree + 1)**dim]
basis_expansion = tf.reduce_prod(samples_centered**grid, -1)
return tf.transpose(basis_expansion)
def resize_and_center_cropped_inputs():
"""Deterministically resize to shorter side and center crop."""
input_shape = tf.compat.v1.shape(inputs)
input_height_t = input_shape[H_AXIS]
input_width_t = input_shape[W_AXIS]
ratio_cond = (input_height_t / input_width_t > (self.height / self.width))
# pylint: disable=g-long-lambda
resized_height = control_flow_util.smart_cond(
ratio_cond,
lambda: tf.cast(self.width * input_height_t / input_width_t,
input_height_t.dtype), lambda: self.height)
resized_width = control_flow_util.smart_cond(
ratio_cond, lambda: self.width,
lambda: tf.cast(self.height * input_width_t / input_height_t,
input_width_t.dtype))
# pylint: enable=g-long-lambda
resized_inputs = tf.image.resize(
images=inputs, size=tf.stack([resized_height, resized_width]))
img_hd_diff = resized_height - self.height
img_wd_diff = resized_width - self.width
bbox_h_start = tf.cast(img_hd_diff / 2, tf.int32)
bbox_w_start = tf.cast(img_wd_diff / 2, tf.int32)
bbox_begin = tf.stack([0, bbox_h_start, bbox_w_start, 0])
bbox_size = tf.stack([-1, self.height, self.width, -1])
outputs = tf.slice(resized_inputs, bbox_begin, bbox_size)
return outputs
def _slice(x, begin, size, axis):
"""Slice wrapper which has XLA-compatible gradients."""
# Not in an XLA context => don't need a custom gradient, can use pad.
if (tf.executing_eagerly()
or not control_flow_util.GraphOrParentsInXlaContext(
tf1.get_default_graph())):
return tf.slice(x, begin, size)
# TODO(b/143313126): Upstream something like this to TF proper.
@tf.custom_gradient
def _custom_slice_helper(x):
"""tf.slice uses tf.pad for backprop, unknown shape breaks XLA."""
def grad(dy):
dshp = tf.shape(x) - tf.shape(dy)
z = tf.zeros(tf.where(tf.equal(0, dshp), tf.shape(x), dshp),
dtype=x.dtype)
dx = tf.roll(tf.concat([dy, z], axis=axis),
shift=begin[axis],
axis=axis)
return dx
return tf.slice(x, begin, size), grad
return _custom_slice_helper(x)
def call(self, inputs):
inputs_shape = tf.compat.v1.shape(inputs)
img_hd = inputs_shape[H_AXIS]
img_wd = inputs_shape[W_AXIS]
img_hd_diff = img_hd - self.target_height
img_wd_diff = img_wd - self.target_width
checks = []
checks.append(
tf.compat.v1.assert_non_negative(
img_hd_diff,
message='The crop height {} should not be greater than input '
'height.'.format(self.target_height)))
checks.append(
tf.compat.v1.assert_non_negative(
img_wd_diff,
message='The crop width {} should not be greater than input '
'width.'.format(self.target_width)))
with tf.control_dependencies(checks):
bbox_h_start = tf.cast(img_hd_diff / 2, tf.int32)
bbox_w_start = tf.cast(img_wd_diff / 2, tf.int32)
bbox_begin = tf.stack([0, bbox_h_start, bbox_w_start, 0])
bbox_size = tf.stack(
[-1, self.target_height, self.target_width, -1])
outputs = tf.slice(inputs, bbox_begin, bbox_size)
return outputs
def relative_attn_bias(rel_bias, num_heads, decode_step=None):
"""Computes attention bias based on relative positions.
Content-based relative position attention bias was used in:
https://arxiv.org/pdf/1803.02155.
Non-content-based relative position attention bias was used in:
https://arxiv.org/abs/1606.01933.
Args:
rel_bias: Relative bias variable of shape [num_heads, 2 * length].
num_heads: Number of attention heads.
decode_step: Optional decode step, used for slicing during decoding.
Returns:
A [..., length, num_heads, length] tensor with queries.
"""
num_rel_pos = rel_bias.shape[-1]
length = num_rel_pos // 2
if tf.is_tensor(decode_step):
# This is decoding so we need to select the current slice within rel_bias.
# E.g.: len_k = 3, decode_step = 1
# We have: rel_bias = [-2, -1, 0, 1, 2, 3]
# We want: [-1, 0, 1]
# We slice at len_k - decode_step - 1 = 1
rel_bias = tf.reshape(rel_bias, [1, num_heads, num_rel_pos])
start = ((length - 1) - decode_step)
rel_bias = tf.slice(rel_bias, [0, 0, start], [1, num_heads, length])
return rel_bias
# Now we have to shift in order to compute relative biases.
# Example: length = 3
# Say we want: [[0, 1, 2], [-1, 0, 1], [-2, -1, 0]]
# Start: [[-2, -1, 0, 1, 2, 3], [-2, -1, 0, 1, 2, 3], [-2, -1, 0, 1, 2, 3]]
# We linearize: [-2, -1, 0, 1, 2, 3, -2, -1, 0, 1, 2, 3, -2, -1, 0, 1, 2, 3]
# We slice: [-2, -1, 0, 1, 2, 3, -2, -1, 0, 1, 2, 3, -2, -1, 0]
# We reshape: [[-2, -1, 0, 1, 2], [3, -2, -1, 0, 1], [2, 3, -2, -1, 0]]
# We slice: [[0, 1, 2], [-1, 0, 1], [-2, -1, 0]]
# Tadaaa!
# [heads, len_q * num_rel_pos]
rel_bias = tf.tile(rel_bias, [1, length])
# [heads, len_q * (num_rel_pos - 1)]
num_rel_pos -= 1
rel_bias = rel_bias[Ellipsis, :length * num_rel_pos]
# [heads, len_q, num_rel_pos - 1]
# Now every row is shifted by 1 to the right.
rel_bias = tf.reshape(rel_bias, [num_heads, length, num_rel_pos])
# [heads, len_q, len_k]
# Slice the overlapping elements from start.
rel_bias = rel_bias[Ellipsis, num_rel_pos - length:]
# [len_q, heads, len_k]
rel_bias = tf.transpose(rel_bias, [1, 0, 2])
return rel_bias
def _cross_suppression(boxes, box_slice, iou_threshold, inner_idx):
batch_size = tf.shape(boxes)[0]
new_slice = tf.slice(boxes, [0, inner_idx * NMS_TILE_SIZE, 0],
[batch_size, NMS_TILE_SIZE, 4])
iou = box_utils.bbox_overlap(new_slice, box_slice)
ret_slice = tf.expand_dims(
tf.cast(tf.reduce_all(iou < iou_threshold, [1]), box_slice.dtype),
2) * box_slice
return boxes, ret_slice, iou_threshold, inner_idx + 1
def _suppression_loop_body(boxes, iou_threshold, output_size, idx):
"""Process boxes in the range [idx*NMS_TILE_SIZE, (idx+1)*NMS_TILE_SIZE).
Args:
boxes: a tensor with a shape of [batch_size, anchors, 4].
iou_threshold: a float representing the threshold for deciding whether boxes
overlap too much with respect to IOU.
output_size: an int32 tensor of size [batch_size]. Representing the number
of selected boxes for each batch.
idx: an integer scalar representing induction variable.
Returns:
boxes: updated boxes.
iou_threshold: pass down iou_threshold to the next iteration.
output_size: the updated output_size.
idx: the updated induction variable.
"""
num_tiles = tf.shape(boxes)[1] // NMS_TILE_SIZE
batch_size = tf.shape(boxes)[0]
# Iterates over tiles that can possibly suppress the current tile.
box_slice = tf.slice(boxes, [0, idx * NMS_TILE_SIZE, 0],
[batch_size, NMS_TILE_SIZE, 4])
_, box_slice, _, _ = tf.while_loop(
lambda _boxes, _box_slice, _threshold, inner_idx: inner_idx < idx,
_cross_suppression, [boxes, box_slice, iou_threshold,
tf.constant(0)])
# Iterates over the current tile to compute self-suppression.
iou = box_utils.bbox_overlap(box_slice, box_slice)
mask = tf.expand_dims(
tf.reshape(tf.range(NMS_TILE_SIZE), [1, -1]) > tf.reshape(
tf.range(NMS_TILE_SIZE), [-1, 1]), 0)
iou *= tf.cast(tf.logical_and(mask, iou >= iou_threshold), iou.dtype)
suppressed_iou, _, _ = tf.while_loop(
lambda _iou, loop_condition, _iou_sum: loop_condition,
_self_suppression,
[iou, tf.constant(True),
tf.reduce_sum(iou, [1, 2])])
suppressed_box = tf.reduce_sum(suppressed_iou, 1) > 0
box_slice *= tf.expand_dims(1.0 - tf.cast(suppressed_box, box_slice.dtype),
2)
# Uses box_slice to update the input boxes.
mask = tf.reshape(tf.cast(tf.equal(tf.range(num_tiles), idx), boxes.dtype),
[1, -1, 1, 1])
boxes = tf.tile(tf.expand_dims(
box_slice, [1]), [1, num_tiles, 1, 1]) * mask + tf.reshape(
boxes, [batch_size, num_tiles, NMS_TILE_SIZE, 4]) * (1 - mask)
boxes = tf.reshape(boxes, [batch_size, -1, 4])
# Updates output_size.
output_size += tf.reduce_sum(
tf.cast(tf.reduce_any(box_slice > 0, [2]), tf.int32), [1])
return boxes, iou_threshold, output_size, idx + 1
def _custom_slice_helper(x):
"""tf.slice uses tf.pad for backprop, unknown shape breaks XLA."""
def grad(dy):
dshp = tf.shape(x) - tf.shape(dy)
z = tf.zeros(tf.where(tf.equal(0, dshp), tf.shape(x), dshp),
dtype=x.dtype)
dx = tf.roll(tf.concat([dy, z], axis=axis),
shift=begin[axis],
axis=axis)
return dx
return tf.slice(x, begin, size), grad
def _symbols_to_logits_fn(target_ids, cache, i):
"""Generate logits for next candidate IDs.
Args:
target_ids: Current decoded sequences. int tensor with shape
[batch_size, i + 1]
cache: dictionary of values storing the encoder output, encoder-decoder
attention bias, and previous decoder attention values.
i: Loop index
Returns:
Tuple of
(logits with shape [batch_size * beam_size, vocab_size],
updated cache values)
"""
decoder_input = tf.slice(
target_ids, [0, tf.maximum(tf.cast(0, i.dtype), i - 1)],
[target_ids.shape[0], 1])
self_attention_mask = tf.slice(decoder_self_attention_mask,
[0, 0, i, 0],
[1, 1, 1, max_decode_len])
# Preprocess decoder input by getting embeddings and adding timing signal.
decoder_input = self.embeder(decoder_input,
1,
start_pos=i,
training=False)
decoder_output = self.decoder(decoder_input,
self_attention_mask,
cache.get("encoder_output"),
cache.get("encoder_mask"),
cache=cache,
decode_i=i,
training=False)
logits = self.embeder.linear(decoder_output)
logits = tf.squeeze(logits, axis=[1])
return logits
def call(self, inputs):
"""Implements call() for the layer."""
if self._use_dynamic_slicing:
input_shape = tf_utils.get_shape_list(inputs, expected_rank=3)
seq_length = input_shape[1]
width = input_shape[2]
position_embeddings = tf.expand_dims(tf.slice(
self._position_embeddings, [0, 0], [seq_length, width]),
axis=0)
else:
position_embeddings = tf.expand_dims(self._position_embeddings,
axis=0)
return position_embeddings
def call(self, x):
shape = x.shape
rank = len(shape)
dim = self.dimension + 1
# Assume 1 batch_dim.
index = [0] * len(self.resolution)
y = x
paddings = np.zeros((rank, 2), dtype=np.int32)
paddings[dim, 0] = 1
y = tf.pad(y, paddings)
rem_dims = rank - 1 - len(index[:dim])
slice_inds = [0] + index[:dim] + [0] * rem_dims
return tf.slice(y, slice_inds, shape)
def random_cropped_inputs():
"""Cropped inputs with stateless random ops."""
input_shape = tf.compat.v1.shape(inputs)
crop_size = tf.stack(
[input_shape[0], self.height, self.width, input_shape[3]])
check = tf.Assert(
tf.reduce_all(input_shape >= crop_size),
[self.height, self.width])
with tf.control_dependencies([check]):
limit = input_shape - crop_size + 1
offset = stateless_random_ops.stateless_random_uniform(
tf.compat.v1.shape(input_shape),
dtype=crop_size.dtype,
maxval=crop_size.dtype.max,
seed=self._rng.make_seeds()[:, 0]) % limit
return tf.slice(inputs, offset, crop_size)
def points_select_first_n(features, num_points, keys=("image", )):
"""Resize the number of points.
Args:
features: Dictionary of data features to preprocess.
num_points: The target number of points.
keys: On which keys to apply this function.
Returns:
Features with resized number of points.
"""
for key in keys:
features[key] = tf.slice(features[key],
begin=(0, 0),
size=(num_points, -1))
return features
def _slice(tensor, dim, start, end):
"""Slices the tensor along given dimension."""
# Performs a slice along the dimension dim. E.g. for tensor t of rank 3,
# _slice(t, 1, 3, 5) is same as t[:, 3:5].
# For a slice unbounded to the right, set end=0: _slice(t, 1, -3, 0) is same
# as t[:, -3:].
rank = len(tensor.shape.as_list())
if start < 0:
start += tf.compat.dimension_value(tensor.shape.as_list()[dim])
if end <= 0:
end += tf.compat.dimension_value(tensor.shape.as_list()[dim])
slice_begin = np.zeros(rank, dtype=np.int32)
slice_begin[dim] = start
slice_size = -np.ones(rank, dtype=np.int32)
slice_size[dim] = end - start
return tf.slice(tensor, slice_begin, slice_size)
def _inverse(self, y):
ndims = prefer_static.rank(y)
shifted_y = tf.pad(
tf.slice(
y, tf.zeros(ndims, dtype=tf.int32),
prefer_static.shape(y) -
tf.one_hot(ndims + self.axis, ndims, dtype=tf.int32)
), # Remove the last entry of y in the chosen dimension.
paddings=tf.one_hot(
tf.one_hot(ndims + self.axis, ndims, on_value=0, off_value=-1),
2,
dtype=tf.int32
) # Insert zeros at the beginning of the chosen dimension.
)
return y - shifted_y
def aligned_random_crop(waves, frame_length):
"""Get aligned random crops from batches of input waves."""
n, t = waves[0].shape
crop_t = frame_length * (t // frame_length - 1)
offsets = [
tf.random.uniform(shape=(),
minval=0,
maxval=t - crop_t,
dtype=tf.int32) for _ in range(n)
]
waves_unbatched = [tf.split(w, n, axis=0) for w in waves]
wave_crops = [[
tf.slice(w, begin=[0, o], size=[1, crop_t])
for w, o in zip(ws, offsets)
] for ws in waves_unbatched]
wave_crops = [tf.concat(wc, axis=0) for wc in wave_crops]
return wave_crops
def body_fn(vecs, i):
# Slice out the vector w.r.t. which we're orthogonalizing the rest.
vecs_ndims = ps.rank(vecs)
select_axis = (ps.range(vecs_ndims) == vecs_ndims - 1)
start = ps.where(select_axis, i, ps.zeros([vecs_ndims], i.dtype))
size = ps.where(select_axis, 1, ps.shape(vecs))
u = tf.math.l2_normalize(tf.slice(vecs, start, size), axis=-2)
# TODO(b/171730305): XLA can't handle this line...
# u = tf.math.l2_normalize(vecs[..., i, tf.newaxis], axis=-2)
# Find weights by dotting the d x 1 against the d x n.
weights = tf.einsum('...dm,...dn->...n', u, vecs)
# Project out vector `u` from the trailing vectors.
masked_weights = tf.where(tf.range(n) > i, weights,
0.)[..., tf.newaxis, :]
vecs = vecs - tf.math.multiply_no_nan(u, masked_weights)
tensorshape_util.set_shape(vecs, vectors.shape)
return vecs, i + 1
def pad_or_trim_to(x, shape, pad_val=0):
"""Pad and slice x to the given shape.
Args:
x: A tensor.
shape: The shape of the returned tensor.
pad_val: An int or float used to pad x.
Returns:
'x' is padded with pad_val and sliced so that the result has the given
shape.
"""
pad = shape - tf.minimum(tf.shape(x), shape)
zeros = tf.zeros_like(pad)
x = tf.pad(x, tf.stack([zeros, pad], axis=1), constant_values=pad_val)
# If dim-i is larger than shape[i], we slice [0:shape[i]] for dim-i.
return tf.reshape(tf.slice(x, zeros, shape), shape)
