add test and numpy/tf implementation for slice_nd

Author: mikel-zhobro

returnn/tf/util/basic.py

@@ -3631,6 +3631,33 @@ def windowed_nd(source, window_size, window_left=None, window_right=None,
     return final
 
 
+def slice_nd2(x, start, size):
+  """
+  This is a more generic slice function, where arbitrary many common axis between x and start are allowed.
+  Here we assume that x and start have their axis layed in the same order.
+
+  :param tf.Tensor x: shape (B, T1, ..., Tn, D)
+  :param tf.Tensor start: shape (B,T1 .., Tn-1), int32 which automatically indicates n as the slice-axis
+  :param int axis: in the range [0..n-1]
+  :param int|tf.Tensor size: scalar
+  :return: ret[b, t1, .., tn-1, 0..size, :] = x[b, t1, .., tn-1, start[B, t1, .., tn-1]+0..size, :]
+    In case the slices go out of bounds of the slice dimension and we will pad with zeros.
+  :rtype: tf.Tensor
+  """
+  with tf.name_scope("slice_nd"):
+    shape = x.shape
+    len_common_dims = len(start.shape)  # nr of common dims
+    slice_dim = shape[len_common_dims]  # dim of axis to be sliced
+    assert size < slice_dim, "Slice size cannot be bigger than the dimension to be sliced."
+    # Create indexes for the slices where slice_idx[B,T1 .., Tn-1] = start[B,T1 .., Tn-1] + range(size)
+    slice_idx = tf.tile(tf.expand_dims(start, -1), [1] * len_common_dims + [size]) + tf.range(size)  # (B,T1 .., Tn-1, size)
+    mask = tf.logical_or(tf.greater(slice_idx, slice_dim - 1), tf.less(slice_idx, 0))  # (B,T1 .., Tn-1, size)
+    slice_idx = tf.clip_by_value(slice_idx, 0, slice_dim - 1)  # cliped slice idx
+    res = tf.gather(x, slice_idx, axis=len_common_dims, batch_dims=len_common_dims)
+    res = where_bc(mask, tf.zeros_like(res), res)  # zero-padding
+    return res
+
+
 def slice_nd(x, start, size):
   """
   :param tf.Tensor x: shape (B, T, ...)

tests/test_TFNetworkRecLayer.py

@@ -3328,7 +3328,7 @@ def test_rec_subnet_simple_rnn():
     print("rnn_cell also fine.")
 
 
-def check_reclayer_optimize_out(subnet_layer_dict, other_subnet_layers=None, shared_base_net=None, rtol=1e-4):
+def check_reclayer_optimize_out(subnet_layer_dict, other_subnet_layers=None, shared_base_net=None, from_=None, rtol=1e-4):
   """
   :param dict[str] subnet_layer_dict: opts for the output layer inside the rec-layer subnet
   :param dict[str,dict[str]] other_subnet_layers: other layers for the rec-layer subnet
@@ -3344,7 +3344,7 @@ def check_reclayer_optimize_out(subnet_layer_dict, other_subnet_layers=None, sha
   subnet_layer_dict.setdefault("from", ["data:source"])
   rec_layer_dict = {
     "class": "rec",
-    "from": ["data"],
+    "from": ["data"] if from_ is None else [from_],
     "unit": {"output": subnet_layer_dict},
     "n_out": n_out,
     "is_output_layer": True
@@ -3598,6 +3598,31 @@ def test_reclayer_optimize_out_access_split():
     other_subnet_layers={"split": {"class": "split", "from": ["data:source"], "size_splits": [5, 8]}})
 
 
+def test_reclayer_optimize_out_slice_nd():
+  def random_start_positions(source, **kwargs):
+    import tensorflow as tf
+    enc = source(0, as_data=True, enforce_batch_major=True, auto_convert=False)
+    enc_shape = tf.shape(enc.placeholder)
+    enc_time_dim = enc_shape[enc.time_dim_axis]
+    return tf.random.uniform(enc_shape[:-1], 0, enc_time_dim-2, dtype=tf.dtypes.int32)
+
+  check_reclayer_optimize_out(
+    {"class": "linear", "activation": None, "from": ["encoder_reduced"]},
+    from_="position",
+    other_subnet_layers={
+      "window": {"class": "slice_nd", "from": "base:encoder", "start": "data:source", "size": None, "min_size": 1, "is_output_layer": True},
+      "encoder_reduced": {"class": "reduce", "mode": "sum", "axis": "T", "from": ["base:encoder"], "is_output_layer": True}
+      },
+    shared_base_net={
+      "encoder": {"class": "copy", "from": "data", "is_output_layer": True},
+      "position": {
+        "class": "eval", "from": "encoder", "is_output_layer": True,
+        "eval": random_start_positions,
+        "out_type": {"batch_dim_axis": 0, "time_dim_axis": 1, "shape": (None,), "sparse": True, "dtype": "int32", "dim": None}}
+      }
+    )
+
+
 def test_reclayer_att_with_kv_in_rec():
   net_dict = {
     'decision': {'class': 'decide', 'from': ['output'], 'loss': 'edit_distance', 'loss_opts': {}, 'target': 'classes'},

tests/test_TFUtil.py

@@ -1637,6 +1637,8 @@ def test_windowed_nd_big():
 
 
 def naive_slice_nd(x, start, size):
+  # old implementation, check out naive_slice_nd2
+
   slices_shape = [x.shape[0], size] + list(x.shape)[2:]
   ys = numpy.zeros(shape=slices_shape)
   for i in range(len(start)):
@@ -1653,6 +1655,63 @@ def naive_slice_nd(x, start, size):
   return ys
 
 
+def naive_slice_nd2(x, start, size):
+  # Assuming that x: [B, T1, T2, .., Tn, D] and start: [B, T1, .., Tn-1]
+  # i.e. the dimensions of x and start are ordered accordingly.
+  # (Otherwise we should require the slice axis too.)
+
+  len_common_dims = len(start.shape)
+  slice_shape = (size,) + x.shape[len_common_dims+1:]
+  result_shape = start.shape[0:len_common_dims] + slice_shape  # shape of output
+  result = numpy.zeros(result_shape)
+
+  slice_axis_dim = x.shape[len_common_dims]  # dim of axis being sliced
+  for index, start_position in numpy.ndenumerate(start):
+    end_position = min(start_position+size, slice_axis_dim)  # padding required
+
+    # no padding
+    padding = ((0,0),)
+    for i in range(1, len(slice_shape)):
+        padding += ((0, 0),)
+
+    # if required replace the first padding tuple, which corresponds to the slice axis
+    if end_position < start_position+size:
+      padding = ((0,size - end_position + start_position),) + padding[1:]
+    result[index] = numpy.pad(x[index][start_position:end_position], padding, mode='constant', constant_values=0)
+  return result
+
+
+def test_slice_nd_multi_dim():
+  n_batch = 2
+  n_time_1 = 2
+  n_time_2 = 3  # slice axis
+  n_dim = 2
+  size = 2
+  source = numpy.arange(24, dtype=numpy.float32).reshape(n_batch, n_time_1, n_time_2, n_dim).astype("float32")
+  start = numpy.array([[0,1],[1,2]]).astype("int32")
+  naive = naive_slice_nd2(source, start, size)
+  source_tf = tf.constant(source)
+  real = slice_nd2(source_tf, start=start, size=size).eval()
+  print("source:")
+  print(source)
+  print("naive:")
+  print(naive)
+  print("real:")
+  print(real)
+  expected_output = numpy.array(
+    [[[[0, 1],
+       [2, 3]],
+      [[8, 9],
+       [10, 11]]],
+
+     [[[14, 15],
+       [16, 17]],
+      [[22, 23],
+       [0, 0]]]])  # padding
+  numpy.testing.assert_almost_equal(naive, expected_output)
+  numpy.testing.assert_almost_equal(real, expected_output)
+
+
 def test_slice_nd_small():
   n_batch = 3
   n_time = 4
@@ -1662,7 +1721,7 @@ def test_slice_nd_small():
   source = numpy.arange(1, n_batch*n_time*n_dim + 1, dtype=numpy.float32).reshape(n_batch, n_time, n_dim).astype("float32")
   source_tf = tf.constant(source)
   naive = naive_slice_nd(source, start, size)
-  real = slice_nd(source_tf, start=start, size=size).eval()
+  real = slice_nd2(source_tf, start=start, size=size).eval()
   print("source:")
   print(source)
   print("naive:")
@@ -1682,7 +1741,7 @@ def test_slice_nd_big():
   source = numpy.random.random((n_batch, n_time, n_dim)).astype("float32")
   source_tf = tf.constant(source)
   naive = naive_slice_nd(source, start, size)
-  real = slice_nd(source_tf, start=start, size=size).eval()
+  real = slice_nd2(source_tf, start=start, size=size).eval()
   print("source:")
   print(source)
   print("naive:")