def compute_output_shape(self, input_shape):
output_shape = self.forward_layer.compute_output_shape(input_shape)
if self.return_state:
state_shape = tf_utils.convert_shapes(
output_shape[1:], to_tuples=False
)
output_shape = tf_utils.convert_shapes(
output_shape[0], to_tuples=False
)
else:
output_shape = tf_utils.convert_shapes(
output_shape, to_tuples=False
)
if self.merge_mode == "concat":
output_shape = output_shape.as_list()
output_shape[-1] *= 2
output_shape = tf.TensorShape(output_shape)
elif self.merge_mode is None:
output_shape = [output_shape, copy.copy(output_shape)]
if self.return_state:
if self.merge_mode is None:
return output_shape + state_shape + copy.copy(state_shape)
return [output_shape] + state_shape + copy.copy(state_shape)
return output_shape
def compute_output_shape(self, input_shape):
if self._output_shape is None:
# Make use of existing autocomputation but provide Lambda-specific
# error message. This is always safe to run even when the outer context
# is Graph mode because Lambda layers don't have side effects such as
# `add_loss`.
with tf.__internal__.eager_context.eager_mode():
try:
return super().compute_output_shape(input_shape)
except NotImplementedError:
raise NotImplementedError(
'We could not automatically infer the shape of the Lambda\'s '
'output. Please specify `output_shape` for this Lambda.'
)
if callable(self._output_shape):
output_shapes = self._output_shape(input_shape)
return tf_utils.convert_shapes(output_shapes, to_tuples=False)
# Output shapes are passed directly and don't include batch dimension.
input_tensor_shape = tf_utils.convert_shapes(input_shape,
to_tuples=False)
batch_size = tf.nest.flatten(
input_tensor_shape)[0][0] if input_shape else None
def _add_batch(shape):
return tf.TensorShape([batch_size] + shape.as_list())
output_shapes = tf_utils.convert_shapes(self._output_shape,
to_tuples=False)
return tf.nest.map_structure(_add_batch, output_shapes)
def compute_output_shape(self, input_shape):
input_shape = tf_utils.convert_shapes(input_shape, to_tuples=False)
child_input_shape = tf.nest.map_structure(self._remove_timesteps, input_shape)
child_output_shape = self.layer.compute_output_shape(child_input_shape)
child_output_shape = tf_utils.convert_shapes(
child_output_shape, to_tuples=False)
timesteps = tf_utils.convert_shapes(input_shape)
timesteps = tf.nest.flatten(timesteps)[1]
def insert_timesteps(dims):
dims = dims.as_list()
return tf.TensorShape([dims[0], timesteps] + dims[1:])
return tf.nest.map_structure(insert_timesteps, child_output_shape)
def build(self, input_shape):
input_shape = tf_utils.convert_shapes(input_shape, to_tuples=False)
input_dims = tf.nest.flatten(
tf.nest.map_structure(lambda x: x.ndims, input_shape))
if any(dim < 3 for dim in input_dims):
raise ValueError(
'`TimeDistributed` Layer should be passed an `input_shape ` '
'with at least 3 dimensions, received: ' + str(input_shape))
# Don't enforce the batch or time dimension.
self.input_spec = tf.nest.map_structure(
lambda x: InputSpec(shape=[None, None] + x.as_list()[2:]), input_shape)
child_input_shape = tf.nest.map_structure(self._remove_timesteps, input_shape)
child_input_shape = tf_utils.convert_shapes(child_input_shape)
super(TimeDistributed, self).build(tuple(child_input_shape))
self.built = True
def compute_mask(self, inputs, mask=None):
"""Computes an output mask tensor for Embedding layer.
This is based on the inputs, mask, and the inner layer.
If batch size is specified:
Simply return the input `mask`. (An rnn-based implementation with
more than one rnn inputs is required but not supported in tf.keras yet.)
Otherwise we call `compute_mask` of the inner layer at each time step.
If the output mask at each time step is not `None`:
(E.g., inner layer is Masking or RNN)
Concatenate all of them and return the concatenation.
If the output mask at each time step is `None` and the input mask is not
`None`:(E.g., inner layer is Dense)
Reduce the input_mask to 2 dimensions and return it.
Otherwise (both the output mask and the input mask are `None`):
(E.g., `mask` is not used at all)
Return `None`.
Arguments:
inputs: Tensor with shape [batch size, timesteps, ...] indicating the
input to TimeDistributed. If static shape information is available for
"batch size", `mask` is returned unmodified.
mask: Either None (indicating no masking) or a Tensor indicating the
input mask for TimeDistributed. The shape can be static or dynamic.
Returns:
Either None (no masking), or a [batch size, timesteps, ...] Tensor with
an output mask for the TimeDistributed layer with the shape beyond the
second dimension being the value of the input mask shape(if the computed
output mask is none), an output mask with the shape beyond the first
dimension being the value of the mask shape(if mask is not None) or
output mask with the shape beyond the first dimension being the
value of the computed output shape.
"""
# cases need to call the layer.compute_mask when input_mask is None:
# Masking layer and Embedding layer with mask_zero
input_shape = tf.nest.map_structure(
lambda x: tf.TensorShape(K.int_shape(x)), inputs)
input_shape = tf_utils.convert_shapes(input_shape, to_tuples=False)
batch_size = tf_utils.convert_shapes(input_shape)
batch_size = tf.nest.flatten(batch_size)[0]
is_ragged_input = tf.nest.map_structure(
lambda x: isinstance(x, tf.RaggedTensor), inputs)
is_ragged_input = generic_utils.to_list(tf.nest.flatten(is_ragged_input))
if batch_size and not self._always_use_reshape or any(is_ragged_input):
# batch size matters, we currently do not handle mask explicitly, or if
# the layer always uses reshape approach, or the input is a ragged tensor.
return mask
inner_mask = mask
if inner_mask is not None:
inner_mask_shape = self._get_shape_tuple((-1,), mask, 2)
inner_mask = K.reshape(inner_mask, inner_mask_shape)
inner_input_shape = tf.nest.map_structure(
lambda tensor: self._get_shape_tuple((-1,), tensor, 2), inputs)
inner_inputs = tf.__internal__.nest.map_structure_up_to(inputs, tf.reshape, inputs,
inner_input_shape)
output_mask = self.layer.compute_mask(inner_inputs, inner_mask)
if output_mask is None:
if mask is None:
return None
# input_mask is not None, and output_mask is None:
# we should return a not-None mask
output_mask = mask
for _ in range(2, len(K.int_shape(mask))):
output_mask = K.any(output_mask, axis=-1)
else:
# output_mask is not None. We need to reshape it
input_length = tf_utils.convert_shapes(input_shape)
input_length = tf.nest.flatten(input_length)[1]
if not input_length:
input_length = tf.nest.map_structure(lambda x: K.shape(x)[1], inputs)
input_length = tf.nest.flatten(input_length)[0]
output_mask_int_shape = K.int_shape(output_mask)
if output_mask_int_shape is None:
# if the output_mask does not have a static shape,
# its shape must be the same as mask's
if mask is not None:
output_mask_int_shape = K.int_shape(mask)
else:
input_shape = generic_utils.to_list(tf.nest.flatten(input_shape))[0]
output_mask_int_shape = K.compute_output_shape(input_shape)[:-1]
output_mask_shape = self._get_shape_tuple(
(-1, input_length), output_mask, 1, output_mask_int_shape[1:])
output_mask = K.reshape(output_mask, output_mask_shape)
return output_mask
def call(self, inputs, training=None, mask=None):
kwargs = {}
if generic_utils.has_arg(self.layer.call, 'training'):
kwargs['training'] = training
input_shape = tf.nest.map_structure(
lambda x: tf.TensorShape(K.int_shape(x)), inputs)
batch_size = tf_utils.convert_shapes(input_shape)
batch_size = tf.nest.flatten(batch_size)[0]
if batch_size and not self._always_use_reshape:
inputs, row_lengths = K.convert_inputs_if_ragged(inputs)
is_ragged_input = row_lengths is not None
input_length = tf_utils.convert_shapes(input_shape)
input_length = tf.nest.flatten(input_length)[1]
# batch size matters, use rnn-based implementation
def step(x, _):
output = self.layer(x, **kwargs)
return output, []
_, outputs, _ = K.rnn(
step,
inputs,
initial_states=[],
input_length=row_lengths[0] if is_ragged_input else input_length,
mask=mask,
unroll=False)
# pylint: disable=g-long-lambda
y = tf.nest.map_structure(
lambda output: K.maybe_convert_to_ragged(is_ragged_input, output,
row_lengths), outputs)
else:
# No batch size specified, therefore the layer will be able
# to process batches of any size.
# We can go with reshape-based implementation for performance.
is_ragged_input = tf.nest.map_structure(
lambda x: isinstance(x, tf.RaggedTensor), inputs)
is_ragged_input = tf.nest.flatten(is_ragged_input)
if all(is_ragged_input):
input_values = tf.nest.map_structure(lambda x: x.values, inputs)
input_row_lenghts = tf.nest.map_structure(
lambda x: x.nested_row_lengths()[0], inputs)
y = self.layer(input_values, **kwargs)
y = tf.nest.map_structure(tf.RaggedTensor.from_row_lengths, y,
input_row_lenghts)
elif any(is_ragged_input):
raise ValueError('All inputs has to be either ragged or not, '
'but not mixed. You passed: {}'.format(inputs))
else:
input_length = tf_utils.convert_shapes(input_shape)
input_length = tf.nest.flatten(input_length)[1]
if not input_length:
input_length = tf.nest.map_structure(lambda x: tf.compat.v1.shape(x)[1],
inputs)
input_length = generic_utils.to_list(tf.nest.flatten(input_length))[0]
inner_input_shape = tf.nest.map_structure(
lambda x: self._get_shape_tuple((-1,), x, 2), inputs)
# Shape: (num_samples * timesteps, ...). And track the
# transformation in self._input_map.
inputs = tf.__internal__.nest.map_structure_up_to(inputs, tf.reshape, inputs,
inner_input_shape)
# (num_samples * timesteps, ...)
if generic_utils.has_arg(self.layer.call, 'mask') and mask is not None:
inner_mask_shape = self._get_shape_tuple((-1,), mask, 2)
kwargs['mask'] = K.reshape(mask, inner_mask_shape)
y = self.layer(inputs, **kwargs)
# Shape: (num_samples, timesteps, ...)
output_shape = self.compute_output_shape(input_shape)
# pylint: disable=g-long-lambda
output_shape = tf.nest.map_structure(
lambda tensor, int_shape: self._get_shape_tuple(
(-1, input_length), tensor, 1, int_shape[2:]), y, output_shape)
y = tf.__internal__.nest.map_structure_up_to(y, tf.reshape, y, output_shape)
if not tf.executing_eagerly():
# Set the static shape for the result since it might be lost during
# array_ops reshape, eg, some `None` dim in the result could be
# inferred.
tf.__internal__.nest.map_structure_up_to(
y, lambda tensor, shape: tensor.set_shape(shape), y,
self.compute_output_shape(input_shape))
return y