update SliceNdLayer

Author: mikel-zhobro

returnn/tf/layers/basic.py

@@ -854,52 +854,79 @@ class SliceNdLayer(_ConcatInputLayer):
   e.g. ``x[start:start + size]``.
   This layers allows a different start slice point for each batch,
   in contrast to :class:`SliceLayer`, and the start is variable.
-  See also :class:`GatherNdLayer`.
+  In case start has more than 1 axis we loop over them.
+  E.g. if start: [B,T] we loop over T and slice for each batch normally.
+  See also :class:`GatherLayer`.
   :class:`PrefixInTimeLayer` can recover the original shape (by zero-padding).
   """
   layer_class = "slice_nd"
   recurrent = True
 
   def __init__(self, start, size, min_size=None, **kwargs):
     """
-    :param LayerBase start:
+    We expect the input to have at least one more axis than start and the rest of the axis in front are the same.
+    In case the input has no extra time axis compared to start, we assume slice_nd is pulled out of a rec layer
+    but the input has stood the same.
+
+    :other LayerBase input_data:  shape [B,T0,..,Tn,D] or [B,T0,..,Tn-1,D]
+    :param LayerBase start:  shape [B,T0,..,Tn-1]
+    :other LayerBase output: shape [B,T0,..,Tn',D]
     :param int|None size: if None, it uses the max possible size, and it becomes a dynamic axis
     :param int|None min_size: if size is None, but we want to have a min-size, set this
     """
     super(SliceNdLayer, self).__init__(**kwargs)
-    from returnn.tf.util.basic import slice_nd, where_bc, expand_multiple_dims, DimensionTag
+    from returnn.tf.util.basic import slice_nd, DimensionTag
+    assert start.output.have_batch_axis() and self.input_data.have_batch_axis()
+    self.start = start
+
     x = self.input_data.copy_as_batch_major()
-    assert x.time_dim_axis == 1, "currently only time-axis==1 supported"
+    start = start.output.copy_as_batch_major()
+
+    # make sure axis of start are in input
+    is_equal_opts = dict(ignore_feature_dim=True, allow_same_spatial_dim=True, broadcast_matches=True)
+    for start_axis in range(start.batch_ndim):
+      assert x.get_dim_tag(start_axis).is_equal(start.get_dim_tag(start_axis), **is_equal_opts)
+
+    # Handle the case when layer is pulled out of rec loop but the input hasn't change
+    if self.optimized_out_of_loop_and_unchanged_input(x, start):
+      # add an axis after the last start axis and tile the input Tn-1 times: [B,T0,..,Tn-1,D] -> [B,T0,..,Tn-1,Tn-1,D]
+      tag = start.get_dim_tag(-1)
+      x = x.copy_add_dim_by_tag(tag, True, start.batch_ndim)  # tiles the input
+
+    start = start.get_placeholder_as_batch_major()
     seq_lens = x.get_sequence_lengths() if x.is_time_axis_dynamic() else None
-    self.start = start
-    assert start.output.have_batch_axis() and start.output.batch_shape == (None,)
-    start = start.output.get_placeholder_as_batch_major()
+    slice_axis = len(list(start.shape)) - 1  # slice_axis w/o batch
     if size is None:
       if seq_lens is None:
-        size = tf.maximum(tf.reduce_max(x.batch_shape[1] - start), 0)
+        size = tf.maximum(tf.reduce_max(x.batch_shape[slice_axis] - start), 0)
       else:
         size = tf.maximum(tf.reduce_max(seq_lens - start), 0)
       if min_size is not None:
         size = tf.maximum(size, min_size)
     self.size = size
-    start = tf.expand_dims(start, axis=1)  # (B, T)
     slices = slice_nd(x.placeholder, start=tf.cast(start, tf.int32), size=size)  # (B,size, ...)
-    if seq_lens is not None:
-      mask = tf.greater_equal(tf.range(size)[None, :] + start, seq_lens[:, None])  # (B,T)
-      mask = expand_multiple_dims(mask, list(range(2, x.batch_ndim)))
-      slices = where_bc(mask, tf.zeros_like(slices), slices)
+
     self.output.size_placeholder = x.size_placeholder.copy()
     if isinstance(size, tf.Tensor):
-      self.output.size_placeholder[0] = tf.maximum(seq_lens - tf.reshape(start, tf.shape(seq_lens)), 0)
+      self.output.size_placeholder[slice_axis] = size
       tag = DimensionTag(
         description="sliced-time:%s" % self.get_absolute_name(),
         kind=DimensionTag.Types.Spatial)
-      tag.set_tag_on_size_tensor(self.output.size_placeholder[0])
+      tag.set_tag_on_size_tensor(self.output.size_placeholder[slice_axis])
     else:
       assert isinstance(size, int)
-      self.output.size_placeholder.pop(0, None)  # static time axis
+      self.output.size_placeholder.pop(slice_axis, None)  # static time axis
+
     self.output.placeholder = slices
 
+  @classmethod
+  def optimized_out_of_loop_and_unchanged_input(cls, input_data, start):
+    """
+    :rtype: bool
+    The idea is to check that the axis after the last common axis is a feature axis instead of spatial.
+    """
+    return input_data.get_dim_tag(start.batch_ndim) == input_data.get_dim_tag(input_data.get_feature_batch_axes()[0])
+
   def get_dep_layers(self):
     """
     :rtype: list[LayerBase]
@@ -916,10 +943,17 @@ def get_out_data_from_opts(cls, name, sources=(), start=None, size=None, **kwarg
     :rtype: Data
     """
     input_data = get_concat_sources_data_template(sources).copy_as_batch_major()
-    if start:
-      input_data.beam = SearchBeam.get_combined_beam(input_data.beam, start.output.beam)
+    start = start.output.copy_as_batch_major()
+    input_data.beam = SearchBeam.get_combined_beam(input_data.beam, start.beam)
+
     in_shape = list(input_data.shape)
-    shape = [size] + in_shape[1:]  # (B, size, ...) (w/o batch)
+    start_shape = list(start.shape)
+    slice_axis = len(start_shape) + 1  # w/o B
+
+    if cls.optimized_out_of_loop_and_unchanged_input(input_data, start):
+      slice_axis -= 1
+    shape = start_shape[:] + [size] + in_shape[slice_axis:]  # (B, T1, .., Tn-1, size, ...) (w/o batch)
+
     out_type = input_data.get_kwargs()
     out_type["name"] = "%s_output" % name
     out_type["shape"] = shape
@@ -1176,6 +1210,143 @@ def transform_config_dict(cls, d, network, get_layer):
       d["position"] = get_layer(d["position"])
 
 
+class SliceNdLayer2(_ConcatInputLayer):
+  """
+  This takes out a slice-range from some axis,
+  e.g. ``x[start:start + size]``.
+  This layers allows a different start slice point for each batch,
+  in contrast to :class:`SliceLayer`, and the start is variable.
+  See also :class:`GatherNdLayer`.
+  :class:`PrefixInTimeLayer` can recover the original shape (by zero-padding).
+
+  Gathers slices on a specified axis from the input layer using indices from a ``position`` layer.
+  If the input is a layer of the shape ``[B,T1,T2,F1]``, and start of shape ``[B,T1]``, and size is S,
+  this will yield output of the shape ``[B,T1,S,F1]`` where
+
+  ``output[b,t1,0:S,f1] = input[b,position[b,t1],T2:T2+S,f1]``
+
+  (if ``D`` is the axis to gather from).
+  In general, all shared axes of the input and the positions will be considered as batch-axes.
+
+  The ``position`` argument can also be an ``int``.
+  In this case, this simply gives ``input[position]`` one the specified ``axis``.
+
+  It's basically a wrapper around ``tf.gather``.
+  It provides the same functionality as the deprecated ``GatherNdLayer``, but is more generic.
+  See also :class:`GatherNdLayer`.
+  """
+  layer_class = "slice_nd2"
+  recurrent = True
+
+  def __init__(self, start, size, min_size=None, **kwargs):
+    """
+    We expect the input to have at least one more axis than start and the rest of the axis in front are the same.
+    In case the input has no extra time axis compared to start, we assume slice_nd is pulled out of a rec layer
+    but the input has stood the same.
+
+    :other LayerBase input_data:  shape [B,T0,..,Tn,D] or [B,T0,..,Tn-1,D]
+    :param LayerBase start:  shape [B,T0,..,Tn-1]
+    :other LayerBase output: shape [B,T0,..,Tn',D]
+    :param int|None size: if None, it uses the max possible size, and it becomes a dynamic axis
+    :param int|None min_size: if size is None, but we want to have a min-size, set this
+    """
+    super(SliceNdLayer2, self).__init__(**kwargs)
+    from returnn.tf.util.basic import slice_nd, DimensionTag
+    assert start.output.have_batch_axis() and self.input_data.have_batch_axis()
+    self.start = start
+
+    input_data = self.input_data.copy_as_batch_major()
+    start = start.output.copy_as_batch_major()
+
+    # make sure axis of start are in input
+    is_equal_opts = dict(ignore_feature_dim=True, allow_same_spatial_dim=True, broadcast_matches=True)
+    for start_axis in range(start.batch_ndim):
+      assert input_data.get_dim_tag(start_axis).is_equal(start.get_dim_tag(start_axis), **is_equal_opts)
+
+    # Handle the case when layer is pulled out of rec loop but the input hasn't change
+    if self.optimized_out_of_loop_and_unchanged_input(input_data, start):
+      # add an axis after the last start axis and tile the input Tn-1 times: [B,T0,..,Tn-1,D] -> [B,T0,..,Tn-1,Tn-1,D]
+      tag = start.get_dim_tag(-1)
+      input_data = input_data.copy_add_dim_by_tag(tag, True, start.batch_ndim)  # tiles the input
+
+    start = start.get_placeholder_as_batch_major()
+    seq_lens = input_data.get_sequence_lengths() if input_data.is_time_axis_dynamic() else None
+    slice_axis = len(list(start.shape)) - 1  # slice_axis w/o batch
+    if size is None:
+      if seq_lens is None:
+        size = tf.maximum(tf.reduce_max(input_data.batch_shape[slice_axis] - start), 0)
+      else:
+        size = tf.maximum(tf.reduce_max(seq_lens - start), 0)
+      if min_size is not None:
+        size = tf.maximum(size, min_size)
+    self.size = size
+
+    self.output.size_placeholder = input_data.size_placeholder.copy()
+    if isinstance(size, tf.Tensor):  # size was None in the beginning
+      self.output.size_placeholder[slice_axis] = size
+      tag = DimensionTag(
+        description="sliced-time:%s" % self.get_absolute_name(),
+        kind=DimensionTag.Types.Spatial)
+      tag.set_tag_on_size_tensor(self.output.size_placeholder[slice_axis])
+    else:
+      assert isinstance(size, int)
+      self.output.size_placeholder.pop(slice_axis, None)  # static time axis
+
+    slices = slice_nd(input_data.placeholder, start=tf.cast(start, tf.int32), size=size)  # (B,size, ...)
+    self.output.placeholder = slices
+
+  @classmethod
+  def optimized_out_of_loop_and_unchanged_input(cls, input_data, start):
+    """
+    :rtype: bool
+    The idea is to check that the axis after the last common axis is a feature axis instead of spatial.
+    """
+    return input_data.get_dim_tag(start.batch_ndim) == input_data.get_dim_tag(input_data.get_feature_batch_axes()[0])
+
+  def get_dep_layers(self):
+    """
+    :rtype: list[LayerBase]
+    """
+    return super(SliceNdLayer2, self).get_dep_layers() + [self.start]
+
+  @classmethod
+  def get_out_data_from_opts(cls, name, sources, start, size=None, **kwargs):
+    """
+    :param str name:
+    :param list[LayerBase] sources:
+    :param LayerBase|None start:
+    :param int|None size:
+    :rtype: Data
+    """
+    input_data = get_concat_sources_data_template(sources).copy_as_batch_major()
+    start = start.output.copy_as_batch_major()
+    input_data.beam = SearchBeam.get_combined_beam(input_data.beam, start.beam)
+
+    in_shape = list(input_data.shape)
+    start_shape = list(start.shape)
+    slice_axis = len(start_shape) + 1  # w/o B
+
+    if cls.optimized_out_of_loop_and_unchanged_input(input_data, start):
+      slice_axis -= 1
+    shape = start_shape[:] + [size] + in_shape[slice_axis:]  # (B, T1, .., Tn-1, size, ...) (w/o batch)
+
+    out_type = input_data.get_kwargs()
+    out_type["name"] = "%s_output" % name
+    out_type["shape"] = shape
+    out_type["batch_dim_axis"] = 0
+    return Data(**out_type)
+
+  @classmethod
+  def transform_config_dict(cls, d, network, get_layer):
+    """
+    :param dict[str] d:
+    :param returnn.tf.network.TFNetwork network:
+    :param get_layer:
+    """
+    super(SliceNdLayer2, cls).transform_config_dict(d, network=network, get_layer=get_layer)
+    d["start"] = get_layer(d["start"])
+
+
 class GatherNdLayer(_ConcatInputLayer):
   """
   Warning: This layer is deprecated, use the more general :class:`GatherLayer` instead.