test_TFNetworkLayer.py

# start: nosetests $this_file --nologcapture
from __future__ import division

import _setup_test_env  # noqa
import tensorflow as tf
from nose.tools import assert_equal, assert_not_equal, assert_is_instance
import unittest
import numpy.testing
from pprint import pprint
from returnn.util import better_exchook
from returnn.config import Config
from returnn.tf.network import *
from returnn.tf.layers.basic import *
import returnn.tf.compat as tf_compat
import returnn.tf.util.basic as tf_util
from returnn.tf.util.data import Dim, SpatialDim, FeatureDim, BatchInfo

print("TF version:", tf.__version__)
print("Numpy version:", numpy.__version__)


@contextlib.contextmanager
def make_scope():
  """
  :rtype: tf.compat.v1.Session
  """
  with tf.Graph().as_default() as graph:
    with tf_compat.v1.Session(graph=graph) as session:
      yield session


def make_feed_dict(data_list, same_time=False, n_batch=3, n_time=7):
  """
  :param list[returnn.tf.util.data.Data]|ExternData data_list:
  :param bool same_time:
  :param int n_batch:
  :param int n_time:
  :rtype: dict[tf.Tensor,numpy.ndarray]
  """
  if isinstance(data_list, ExternData):
    data_list = [value for (key, value) in sorted(data_list.data.items())]
  assert n_time > 0 and n_batch > 0
  rnd = numpy.random.RandomState(42)
  existing_sizes = {}  # type: typing.Dict[tf.Tensor,int]
  d = {}
  batch_info = None
  for data in data_list:
    if data.batch and not batch_info:
      batch_info = data.batch
    shape = list(data.batch_shape)
    if data.batch_dim_axis is not None:
      shape[data.batch_dim_axis] = n_batch
    for axis, dim in enumerate(shape):
      if dim is None:
        axis_wo_b = data.get_batch_axis_excluding_batch(axis)
        assert axis_wo_b in data.size_placeholder
        dyn_size = data.size_placeholder[axis_wo_b]
        if dyn_size in existing_sizes:
          shape[axis] = existing_sizes[dyn_size]
          continue
        existing_sizes[dyn_size] = n_time
        shape[axis] = n_time
        dyn_size_v = numpy.array([n_time, max(n_time - 2, 1), max(n_time - 3, 1)])
        if dyn_size_v.shape[0] > n_batch:
          dyn_size_v = dyn_size_v[:n_batch]
        elif dyn_size_v.shape[0] < n_batch:
          dyn_size_v = numpy.concatenate(
            [dyn_size_v, rnd.randint(1, n_time + 1, size=(n_batch - dyn_size_v.shape[0],))], axis=0)
        d[dyn_size] = dyn_size_v
        if not same_time:
          n_time += 1
    print("%r %r: shape %r" % (data, data.placeholder, shape))
    if data.sparse:
      d[data.placeholder] = rnd.randint(0, data.dim or 13, size=shape, dtype=data.dtype)
    else:
      d[data.placeholder] = rnd.normal(size=shape).astype(data.dtype)
  if batch_info:
    batch_dim = batch_info.dim
    if isinstance(batch_dim, int):
      assert batch_dim == n_batch, "invalid batch info %r" % batch_info
    else:
      assert isinstance(batch_dim, tf.Tensor)
      d[batch_dim] = n_batch
  return d


def test_ExternData_init_from_config():
  config = Config({
    "extern_data": {"data": {"dim": 42}},
  })
  extern_data = ExternData()
  with make_scope() as session:
    extern_data.init_from_config(config)
    data = extern_data.data["data"]
    assert data.batch_shape == (None, None, 42)
    assert (data.batch_dim_axis, data.time_dim_axis, data.feature_dim_axis) == (0, 1, 2)


def test_ExternData_init_from_config_dim_none():
  config = Config({
    "extern_data": {"data": {"dim": None}},
  })
  extern_data = ExternData()
  with make_scope() as session:
    extern_data.init_from_config(config)
    data = extern_data.data["data"]
    assert data.batch_shape == (None, None, None)
    assert (data.batch_dim_axis, data.time_dim_axis, data.feature_dim_axis) == (0, 1, 2)


def test_ExternData_init_twice_existing_dim_tags():
  from returnn.tf.util.data import batch_dim
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("feature", dimension=10)
  config = Config({
    "extern_data": {
      "data": {"dim_tags": [batch_dim, time_dim, feat_dim]}  # [B,T,D]
    }
  })
  for _ in range(2):
    with make_scope() as session:
      net = TFNetwork(config=config)
      net.construct_from_dict({"output": {"class": "softmax_over_spatial", "from": "data"}})
      session.run(net.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_LinearLayer():
  from returnn.tf.util.data import batch_dim
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("feature", dimension=5)
  config = Config({
    "extern_data": {
      "data": {"dim_tags": [batch_dim, time_dim, feat_dim]}  # [B,T,D]
    }
  })
  for _ in range(2):
    with make_scope() as session:
      net = TFNetwork(config=config)
      net.construct_from_dict({"output": {"class": "linear", "from": "data", "n_out": 3}})
      session.run(tf_compat.v1.global_variables_initializer())
      session.run(net.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_LinearLayer_in_dim_spatial():
  from returnn.tf.util.data import batch_dim
  time_dim = SpatialDim("time")
  static_spatial_dim = FeatureDim("static-spatial", dimension=3)
  feat_dim = FeatureDim("in-feature", dimension=5)
  out_dim = FeatureDim("out-feature", dimension=7)
  config = Config({
    "extern_data": {
      "data": {"dim_tags": [batch_dim, time_dim, static_spatial_dim, feat_dim]}  # [B,T,D1,D2]
    }
  })
  for _ in range(2):
    with make_scope() as session:
      net = TFNetwork(config=config)
      net.construct_from_dict({
        "output": {"class": "linear", "from": "data", "in_dim": static_spatial_dim, "out_dim": out_dim}})
      layer = net.get_default_output_layer()
      print("Output:", layer.output)
      assert layer.output.dim_tags_set_implicit == {batch_dim, time_dim, out_dim, feat_dim}
      param = layer.params["W"]
      assert isinstance(param, tf.Variable)
      assert param.shape.as_list() == [static_spatial_dim.dimension, out_dim.dimension]
      session.run(tf_compat.v1.global_variables_initializer())
      session.run(layer.output.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_LinearLayer_two_time_dims_allow_broadcast_all_sources():
  from returnn.tf.util.data import batch_dim
  with make_scope() as session:
    time1_dim = SpatialDim("time1")
    time2_dim = SpatialDim("time2")
    feat_dim = FeatureDim("feature", dimension=5)
    out_dim = FeatureDim("feature", dimension=3)
    config = Config({
      "extern_data": {
        "in1": {"dim_tags": [batch_dim, time1_dim, feat_dim]},
        "in2": {"dim_tags": [batch_dim, time2_dim, feat_dim]},
      },
    })
    network = TFNetwork(config=config)
    try:
      network.construct_from_dict({
        "output": {"class": "linear", "from": ["data:in1", "data:in2"], "n_out": 3}})
    except Exception as exc:
      # https://github.com/rwth-i6/returnn/issues/691
      print("Expected exception:", exc)
      assert "require broadcasting" in str(exc)
    else:
      raise Exception(
        "Expect allow_broadcast_all_sources exception, but layer constructed: %s" % network.get_default_output_layer())
    network.construct_from_dict({
      "output": {
        "class": "linear", "from": ["data:in1", "data:in2"], "out_dim": out_dim,
        "out_shape": {batch_dim, time1_dim, time2_dim, out_dim}}})
    output = network.get_default_output_layer().output
    assert output.shape == (None, None, 3)
    session.run(tf_compat.v1.global_variables_initializer())
    session.run(fetches=output.placeholder, feed_dict=make_feed_dict(network.extern_data))


def test_LinearLayer_generic_dim_tags():
  from returnn.tf.util.data import batch_dim, any_feature_dim, any_spatial_dim
  with make_scope() as session:
    time1_dim = SpatialDim("time1")
    time2_dim = SpatialDim("time2", dimension=7)
    feat_dim = FeatureDim("feature", dimension=5)
    out_dim = FeatureDim("feature", dimension=3)
    config = Config({
      "extern_data": {
        "in1": {"dim_tags": [batch_dim, time1_dim, time2_dim, feat_dim]},
        "in2": {"dim_tags": [batch_dim, time2_dim, feat_dim]},
      },
    })
    network = TFNetwork(config=config)
    network.construct_from_dict({
      "output1": {
        "class": "linear", "from": "data:in1", "in_dim": any_feature_dim, "out_dim": out_dim,
        "out_shape": {batch_dim, time1_dim, time2_dim, out_dim},
        "is_output_layer": True}})
    network.construct_from_dict({
      "output2": {
        "class": "linear", "from": "data:in1", "in_dim": time2_dim, "out_dim": out_dim,
        "out_shape": {batch_dim, time1_dim, out_dim, feat_dim},
        "is_output_layer": True}})
    try:
      network.construct_from_dict({
        "output3": {
          "class": "linear", "from": "output2", "in_dim": any_feature_dim, "out_dim": out_dim,
          "is_output_layer": True}})
    except Exception as exc:
      print("Expected exception:", exc)
      assert "not found or unique in input" in str(exc)
    else:
      raise Exception("No exception")
    network.construct_from_dict({
      "output4": {
        "class": "linear", "from": "data:in2", "in_dim": any_spatial_dim, "out_dim": out_dim,
        "out_shape": {batch_dim, out_dim, feat_dim},
        "is_output_layer": True}})
    session.run(tf_compat.v1.global_variables_initializer())
    session.run(
      fetches=[layer.output.placeholder for layer in network.get_output_layers()],
      feed_dict=make_feed_dict(network.extern_data))


def test_LinearLayer_reuse_params_layer_output():
  from returnn.tf.util.data import batch_dim
  with make_scope() as session:
    time_dim = SpatialDim("time")
    data_feat_dim = FeatureDim("feature", dimension=5)
    out_feat_dim = FeatureDim("feature", dimension=7)
    config = Config({
      "extern_data": {
        "data": {"dim_tags": [batch_dim, time_dim, data_feat_dim]},
      },
    })
    network = TFNetwork(config=config)
    network.construct_from_dict({
      "weights": {"class": "variable", "shape": [data_feat_dim, out_feat_dim]},
      "bias": {"class": "variable", "shape": [out_feat_dim]},
      "out1": {
        "class": "linear", "from": "data", "out_dim": out_feat_dim,
        "is_output_layer": True,
        "reuse_params": {
          "map": {
            "W": {"layer_output": "weights"},
            "b": {"layer_output": "bias"}}}},
      "out2_": {
        "class": "dot", "from": ["data", "weights"],
        "red1": data_feat_dim, "red2": data_feat_dim,
        "var1": [batch_dim, time_dim], "var2": out_feat_dim},
      "out2": {
        "class": "combine", "kind": "add", "from": ["out2_", "bias"],
        "is_output_layer": True}
    })
    out1 = network.get_layer("out1").output
    out2 = network.get_layer("out2").output
    params = network.get_params_list()
    assert len(params) == 2  # weights and bias
    session.run(tf_compat.v1.global_variables_initializer())
    out1_np, out2_np = session.run(
      fetches=(out1.placeholder, out2.placeholder), feed_dict=make_feed_dict(network.extern_data))
    numpy.testing.assert_array_equal(out1_np, out2_np)


def test_PadLayer_time():
  n_batch, n_time, n_in = 7, 3, 20
  config = Config({
    "extern_data": {"data": {"dim": n_in}},
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    padding = (2, 3)
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {"class": "pad", "axes": "T", "padding": padding, "mode": "replication", "from": "data:data"}
    })
    out_t = net.get_default_output_layer().output.placeholder
    assert out_t.shape.as_list() == [None, None, n_in]
    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_time, n_in))
    out_v = session.run(out_t, feed_dict={net.extern_data.data["data"].placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, n_time + sum(padding), n_in)
    assert (out_v[:, 0, :] == out_v[:, padding[0], :]).all()
    assert (out_v[:, -1, :] == out_v[:, -1 - padding[1], :]).all()
    numpy.testing.assert_array_equal(in_v, out_v[:, padding[0]:(-padding[1] or None), :])
    # check padding on left
    if padding[0] > 0:
      padded_left_ref = numpy.resize(in_v[:, 0, :], (padding[0], n_batch, n_in)).transpose(1, 0, 2)
      numpy.testing.assert_array_equal(padded_left_ref, out_v[:, :padding[0], :])
    # check padding on right
    if padding[1] > 0:
      padded_right_ref = numpy.resize(in_v[:, -1, :], (padding[1], n_batch, n_in)).transpose(1, 0, 2)
      numpy.testing.assert_array_equal(padded_right_ref, out_v[:, -padding[1]:, :])


def test_PadLayer_feature():
  n_batch, n_time, n_in = 7, 3, 20
  config = Config({
    "extern_data": {"data": {"dim": None}},
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    padding = (2, 3)
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {"class": "pad", "axes": "F", "padding": padding, "mode": "replication", "from": "data:data"}
    })
    out_t = net.get_default_output_layer().output.placeholder
    assert out_t.shape.as_list() == [None, None, None]
    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_time, n_in))
    out_v = session.run(out_t, feed_dict={net.extern_data.data["data"].placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, n_time, n_in + sum(padding))
    assert (out_v[:, :, 0] == out_v[:, :, padding[0]]).all()
    assert (out_v[:, :, -1] == out_v[:, :, -1 - padding[1]]).all()
    numpy.testing.assert_array_equal(in_v, out_v[:, :,  padding[0]:(-padding[1] or None)])
    # check padding on left
    if padding[0] > 0:
      padded_left_ref = numpy.resize(in_v[:, :, 0], (n_batch, n_time, 1))
      numpy.testing.assert_array_equal(padded_left_ref - out_v[:, :, :padding[0]], 0)
    # check padding on right
    if padding[1] > 0:
      padded_left_ref = numpy.resize(in_v[:, :, -1], (n_batch, n_time, 1))
      numpy.testing.assert_array_equal(padded_left_ref - out_v[:, :, -padding[1]:], 0)


def test_PadLayer_no_op():
  # https://github.com/rwth-i6/returnn/issues/687
  n_batch, n_time, n_in = 7, 3, 5
  config = Config({
    "extern_data": {"data": {"shape": (n_in, None)}},  # [B,D,T]
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {'class': 'pad', 'mode': 'constant', 'axes': 'spatial', 'padding': [(0, 0)], 'from': 'data', 'value': 0}
    })
    out = net.get_default_output_layer().output
    out_t = out.placeholder
    assert out_t.shape.as_list() == [None, n_in, None]
    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_in, n_time))
    out_v = session.run(out_t, feed_dict={net.extern_data.data["data"].placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, n_in, n_time)
    numpy.testing.assert_array_equal(in_v, out_v)


def test_PadLayer_window():
  # https://github.com/rwth-i6/returnn/issues/1224
  from returnn.config import Config
  from returnn.tf.engine import Engine
  from returnn.datasets import init_dataset
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim

  time_dim = SpatialDim('time')
  in_dim = FeatureDim('in', 3)
  out_dim = FeatureDim('out', 4)

  def _config_get_network(epoch, **_kwargs):
    window_dim = SpatialDim('window', 3)
    time_dim_ = (window_dim // 2) + time_dim + window_dim.ceildiv_right(2) + (-1)
    flat_dim = window_dim * time_dim_
    flat_dim_ = flat_dim + window_dim
    time_window_dim = time_dim + window_dim

    # This is what PadLayer.get_out_data_from_opts() does.
    # In returnn-common, we would execute that.
    # This should be fine. But this triggers the bug.
    time_dim__ = 1 + time_dim + 1
    time_dim__.declare_same_as(time_dim_)
    flat_dim__ = 0 + flat_dim + 3
    flat_dim__.declare_same_as(flat_dim_)

    net_dict = {
      "#epoch": epoch,  # trigger reinit
      'window': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'pad': {
            'class': 'pad',
            'from': 'base:data:data',
            'axes': time_dim,
            'padding': (1, 1),
            'out_dims': time_dim_,
            'out_shape': {batch_dim, in_dim, time_dim_}
          },
          'expand_dim': {
            'class': 'expand_dims',
            'from': 'pad',
            'axis': 'spatial',
            'dim': window_dim,
            'out_shape': {batch_dim, in_dim, window_dim, time_dim_}
          },
          'merge_dims': {
            'class': 'merge_dims',
            'from': 'expand_dim',
            'axes': (
              window_dim,
              time_dim_
            ),
            'out_dim': flat_dim,
            'out_shape': {batch_dim, in_dim, flat_dim}
          },
          'pad_0': {
            'class': 'pad',
            'from': 'merge_dims',
            'axes': flat_dim,
            'padding': (0, 3),
            'out_dims': flat_dim_,
            'out_shape': {batch_dim, in_dim, flat_dim_}
          },
          'reshape': {
            'class': 'reshape',
            'from': 'pad_0',
            'in_dims': [
              flat_dim_
            ],
            'out_dims': [
              window_dim,
              time_window_dim
            ],
            'extra_deps': ['base:data:data'],
            'out_shape': {batch_dim, in_dim, window_dim, time_window_dim}
          },
          'slice_nd': {
            'class': 'slice_nd',
            'from': 'reshape',
            'size': time_dim,
            'axis': time_window_dim,
            'out_spatial_dim': time_dim,
            'out_shape': {batch_dim, time_dim, in_dim, window_dim}
          },
          'output': {
            'class': 'copy',
            'from': 'slice_nd',
            'out_shape': {batch_dim, time_dim, in_dim, window_dim}
          }
        }
      },
      'reduce': {
        'class': 'reduce',
        'from': 'window',
        'mode': 'mean',
        'axis': (window_dim, in_dim),
        'out_shape': {batch_dim, time_dim}
      },
      'add': {
        'class': 'combine',
        'from': ['dot', 'reduce'],
        'kind': 'add',
        'is_output_layer': True,
        'out_shape': {batch_dim, time_dim, out_dim}
      },
      'reduce_0': {
        'class': 'reduce',
        'from': 'add',
        'mode': 'mean',
        'axis': out_dim,
        'out_shape': {batch_dim, time_dim}
      },
      'dummy': {
        'class': 'copy',
        'from': 'reduce_0',
        'loss': 'as_is',
        'out_shape': {batch_dim, time_dim}
      },
      'weight': {
        'class': 'variable',
        'shape': [
          in_dim,
          out_dim
        ],
        'param_name': 'param',
      },
      'dot': {
        'class': 'dot',
        'from': ['data:data', 'weight'],
        'reduce': in_dim,
        'out_shape': {batch_dim, time_dim, out_dim}
      },
    }
    return net_dict

  config = Config({
    "task": "train", "num_epochs": 2, "start_epoch": 1,
    "get_network": _config_get_network,
    "extern_data": {"data": {"dim_tags": (batch_dim, time_dim, in_dim)}},
  })
  train_dataset = init_dataset(
    {"class": "DummyDataset", "input_dim": in_dim.dimension, "output_dim": 5, "num_seqs": 3})
  engine = Engine(config)
  engine.init_train_from_config(config, train_data=train_dataset)
  engine.train()


def test_concat_sources():
  with make_scope() as session:
    network = TFNetwork(train_flag=True, extern_data=ExternData())
    n_batch = 5
    n_time = 3
    size_placeholder = {0: tf.constant(n_time, dtype=tf.int32, shape=(n_batch,))}
    src1 = InternalLayer(
      name="src1", network=network,
      output=Data(
        name="src1_output", shape=(None, 11), placeholder=tf.zeros((n_batch, n_time, 11)),
        size_placeholder=size_placeholder))
    print("src1 output:", src1.output)
    src2 = InternalLayer(
      name="src2", network=network,
      output=Data(
        name="src2_output", shape=(None, 13), placeholder=tf.zeros((n_batch, n_time, 13)),
        size_placeholder=size_placeholder))
    print("src2 output:", src2.output)
    out_kwargs = dict(name="out", sources=[src1, src2], network=network)
    out_output = CopyLayer.get_out_data_from_opts(**out_kwargs)
    print("out output:", out_output)
    assert out_output.dim == 11 + 13
    out = CopyLayer(output=out_output, **out_kwargs)
    session.run(out.output.placeholder)


def test_concat_sources_batch_dim():
  with make_scope() as session:
    network = TFNetwork(train_flag=True, extern_data=ExternData())
    n_batch = 5
    n_time = 3
    size_placeholder = {0: tf.constant(n_time, dtype=tf.int32, shape=(n_batch,))}
    src1 = InternalLayer(
      name="src1", network=network,
      output=Data(
        name="src1_output", shape=(None, 11), placeholder=tf.zeros((n_batch, n_time, 11)),
        size_placeholder=size_placeholder))
    print("src1 output:", src1.output)
    src2 = InternalLayer(
      name="src2", network=network,
      output=Data(
        name="src2_output", shape=(None, 13), time_dim_axis=0, batch_dim_axis=1,
        placeholder=tf.zeros((n_time, n_batch, 13)),
        size_placeholder=size_placeholder))
    print("src2 output:", src2.output)
    out_kwargs = dict(name="out", sources=[src1, src2], network=network)
    out_output = CopyLayer.get_out_data_from_opts(**out_kwargs)
    print("out output:", out_output)
    assert out_output.dim == 11 + 13
    assert out_output.batch_dim_axis == 0 and out_output.time_dim_axis == 1
    out = CopyLayer(output=out_output, **out_kwargs)
    session.run(out.output.placeholder)


def test_concat_sources_missing_dim():
  with make_scope() as session:
    network = TFNetwork(train_flag=True, extern_data=ExternData())
    n_batch = 5
    n_time = 3
    size_placeholder = {0: tf.constant(n_time, dtype=tf.int32, shape=(n_batch,))}
    src1 = InternalLayer(
      name="src1", network=network,
      output=Data(
        name="src1_output", shape=(None, 11), placeholder=tf.zeros((n_batch, n_time, 11)),
        size_placeholder=size_placeholder))
    print("src1 output:", src1.output)
    src2 = InternalLayer(
      name="src2", network=network,
      output=Data(
        name="src2_output", shape=(13,), time_dim_axis=None, batch_dim_axis=0,
        placeholder=tf.zeros((n_batch, 13)),
        size_placeholder={}))
    print("src2 output:", src2.output)
    out_kwargs = dict(name="out", sources=[src1, src2], network=network)
    out_output = CopyLayer.get_out_data_from_opts(**out_kwargs)
    print("out output:", out_output)
    assert out_output.dim == 11 + 13
    assert out_output.batch_dim_axis == 0 and out_output.time_dim_axis == 1
    out = CopyLayer(output=out_output, **out_kwargs)
    session.run(out.output.placeholder)


def test_concat_sources_dim1():
  with make_scope() as session:
    net_dict = {
      "lin1": {"class": "linear", "activation": "sigmoid", "n_out": 5, "from": "data:data"},
      "lin2": {"class": "linear", "activation": "sigmoid", "n_out": 1, "from": "data:data"},
      "concat": {"class": "copy", "from": ["lin1", "lin2"]},
      "output": {"class": "softmax", "loss": "ce", "from": "concat"}
    }
    config = Config({"debug_print_layer_output_template": True})
    config.update(dict(num_inputs=4, num_outputs=9))
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    assert_equal(network.get_layer("concat").output.shape, (None, 6))
    out = network.get_default_output_layer()
    assert out.output.shape == (None, 9)
    feed_dict = make_feed_dict(network.extern_data.data.values(), same_time=True)
    session.run(tf_compat.v1.global_variables_initializer())
    session.run(out.output.placeholder, feed_dict=feed_dict)


def test_concat_new_dim_tag():
  from returnn.tf.util.data import Dim
  with make_scope():
    n_out = 5
    time_tag = Dim(Dim.Types.Spatial, "time")
    new_time_tag = Dim(Dim.Types.Spatial, "new-time")
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out, "same_dim_tags_as": {"t": time_tag}},
        "classes": {"dim": n_out, "sparse": True, "same_dim_tags_as": {"t": time_tag}}
      }})
    net = TFNetwork(config=config, search_flag=True)
    net.construct_from_dict({
      "data_new": {
        "class": "reinterpret_data", "from": "data",
        "set_dim_tags": {"t": new_time_tag}},
      "output": {"class": "rec", "from": "data", "unit": {
        "prev_out0": {
          "class": "reinterpret_data", "from": "prev:output", "set_sparse": False},
        "prev_out1": {"class": "cast", "from": "prev_out0", "dtype": "float32"},
        "prev_out": {"class": "expand_dims", "from": "prev_out1", "axis": "f"},
        "data_concat": {
          "class": "copy", "from": ["base:data_new", "prev_out"]
        },
        "data_red": {"class": "reduce", "from": "data_concat", "axis": "stag:new-time", "mode": "max"},
        "output_prob": {"class": "softmax", "from": "data_red", "target": "classes", "loss": "ce"},
        "output": {
          "class": "choice", "from": "output_prob", "beam_size": 3, "target": "classes",
          "input_type": "prob", "initial_output": 0}
      }}
    })


def test_ConcatLayer():
  with make_scope() as session:
    net_dict = {
      "lin1": {"class": "linear", "activation": "sigmoid", "n_out": 5, "from": "data:data"},
      "lin2": {"class": "linear", "activation": "sigmoid", "n_out": 3, "from": "data:data"},
      "output": {"class": "concat", "from": [("lin1", "F"), ("lin2", "F"), ("data", "F")]},
    }
    config = Config({"extern_data": {"data": {"dim": 2}}})
    network = TFNetwork(config=config)
    network.construct_from_dict(net_dict)
    out = network.get_default_output_layer()
    assert_equal(out.output.shape, (None, 10))
    feed_dict = make_feed_dict(network.extern_data, same_time=True)
    session.run(tf_compat.v1.global_variables_initializer())
    session.run(out.output.placeholder, feed_dict=feed_dict)


def test_ConcatLayer_range_dyn():
  with make_scope() as session:
    net_dict = {
      "range": {"class": "range_in_axis", "from": "data:data", "axis": "T"},
      "output": {"class": "concat", "from": [("range", "T"), ("range", "T")]},
    }
    config = Config({"extern_data": {"data": {"dim": 2}}})
    network = TFNetwork(config=config)
    network.construct_from_dict(net_dict)
    out = network.get_default_output_layer().output
    assert_equal(out.batch_shape, (None,))
    feed_dict = make_feed_dict(network.extern_data, n_time=7)
    session.run(out.placeholder, feed_dict=feed_dict)
    assert_equal(session.run(out.dim_tags[0].get_dim_value(), feed_dict=feed_dict), 14)


def test_LinearLayer_batch_feature_major():
  with make_scope() as session:
    network = TFNetwork(config=Config(), extern_data=ExternData(), train_flag=True)
    n_in = 3
    n_out = 7
    source = InternalLayer(
      name="source", network=network, output=Data(
        name="source", shape=(n_in, None), time_dim_axis=2, auto_create_placeholders=True))
    assert source.output.feature_dim_axis == 1
    assert source.output.is_batch_feature_major
    out_template = LinearLayer.get_out_data_from_opts(
      name="lin", network=network, n_out=n_out, activation=None, sources=[source])
    out_template.sanity_check()
    assert out_template.shape == (n_out, None) and (out_template.feature_dim_axis, out_template.time_dim_axis) == (1, 2)
    assert out_template.is_batch_feature_major
    with tf_compat.v1.variable_scope("lin"):
      layer = LinearLayer(
        name="lin", network=network, n_out=n_out, activation=None, sources=[source], output=out_template)
    layer.output.sanity_check()
    n_batch = 5
    n_times = [13, 13, 11, 7, 5]
    assert len(n_times) == n_batch
    n_time = max(n_times)
    feed_dict = {
      source.output.placeholder: numpy.random.normal(size=(n_batch, n_in, n_time)).astype("float32"),
      source.output.size_placeholder[1]: numpy.array(n_times, dtype="int32")}
    session.run(tf_compat.v1.global_variables_initializer())
    session.run(layer.output.placeholder, feed_dict=feed_dict)


def test_batch_norm_vars():
  with make_scope() as session:
    n_in, n_out = 2, 3
    config = Config()
    layer_name = "layer1"
    config.update({
      "num_outputs": n_out,
      "num_inputs": n_in,
      "network": {
        layer_name: {
          "class": "linear", "activation": "relu", "batch_norm": {"masked_time": True},
          "n_out": n_out, "is_output_layer": True,
          "from": "data:data"}
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    layer = network.layers[layer_name]
    print("layer:", layer)
    print("layer vars:")
    pprint(layer.params)
    assert layer.use_batch_norm
    bn_prefix = "batch_norm/v2_"
    assert_equal(set(layer.params.keys()), {
      "W", "b", bn_prefix + "beta", bn_prefix + "mean", bn_prefix + "gamma", bn_prefix + "variance"})
    assert_equal(layer.params["W"].get_shape().as_list(), [n_in, n_out])
    assert_equal(layer.params["b"].get_shape().as_list(), [n_out])
    assert_equal(layer.params[bn_prefix + "beta"].get_shape().as_list(), [n_out])
    assert_equal(layer.params[bn_prefix + "gamma"].get_shape().as_list(), [n_out])
    assert_equal(layer.params[bn_prefix + "mean"].get_shape().as_list(), [n_out])
    assert_equal(layer.params[bn_prefix + "variance"].get_shape().as_list(), [n_out])


def _test_batch_norm_param_old_to_new_import(old_version, new_version):
  import tempfile
  model_tmp_dir = tempfile.mkdtemp("tmp-checkpoint")
  model_filename = model_tmp_dir + "/model"
  layer_name = "layer1"
  n_in = 3

  def _make_net_dict(param_version):
    return {
      layer_name: {
        "class": "batch_norm", "from": "data:data", "is_output_layer": True,
        "param_version": param_version, "masked_time": True}
    }

  with make_scope() as session:
    config = Config({"extern_data": {"data": {"dim": n_in}}})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(_make_net_dict(param_version=old_version))
    network.initialize_params(session)
    network.save_params_to_file(filename=model_filename, session=session)
    out_ref = session.run(
      network.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(network.extern_data))
    assert isinstance(out_ref, numpy.ndarray)
    assert not numpy.allclose(out_ref, 0.0)

  with make_scope() as session:
    config = Config({"extern_data": {"data": {"dim": n_in}}})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(_make_net_dict(param_version=new_version))
    network.load_params_from_file(filename=model_filename, session=session)
    out_new = session.run(
      network.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(network.extern_data))
    assert isinstance(out_new, numpy.ndarray)
    assert not numpy.allclose(out_new, 0.0)
    numpy.testing.assert_allclose(out_ref, out_new)


def test_batch_norm_param_v0_to_v1_import():
  _test_batch_norm_param_old_to_new_import(old_version=0, new_version=1)


def test_batch_norm_param_v0_to_v2_import():
  _test_batch_norm_param_old_to_new_import(old_version=0, new_version=2)


def test_batch_norm_param_v1_to_v2_import():
  _test_batch_norm_param_old_to_new_import(old_version=1, new_version=2)


def test_batch_norm_fused():
  n_in = 3
  net_dict = {
    "output": {
      "class": "batch_norm", "from": "data:data",
      "masked_time": False,
      "param_version": 2,
    }
  }

  def _find_fused_bn_op(session_):
    for op in session_.graph.get_operations():
      assert isinstance(op, tf.Operation)
      if "FusedBatchNorm" in op.type:
        return op

  with make_scope() as session:
    config = Config({"extern_data": {"data": {"dim": n_in}}})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    out = network.get_default_output_layer().output.placeholder
    assert _find_fused_bn_op(session)
    network.initialize_params(session)
    out_np = session.run(out, feed_dict=make_feed_dict(network.extern_data))
    assert isinstance(out_np, numpy.ndarray)
    assert not numpy.allclose(out_np, 0.0)

  with make_scope() as session:
    config = Config({"extern_data": {"data": {"dim": n_in}}})
    network = TFNetwork(config=config, train_flag=False)
    network.construct_from_dict(net_dict)
    out = network.get_default_output_layer().output.placeholder
    assert _find_fused_bn_op(session)
    network.initialize_params(session)
    out_np = session.run(out, feed_dict=make_feed_dict(network.extern_data))
    assert isinstance(out_np, numpy.ndarray)
    assert not numpy.allclose(out_np, 0.0)


def test_batch_norm():
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData(), train_flag=True)
    with tf_compat.v1.variable_scope("src_nchw"):
      src_nhwc = InternalLayer(
        name="src_nchw", network=net,
        output=Data(**{
          "name": "src_nchw_output",
          "dim": 16,
          "shape": (None, 16, 16),
          "batch_dim_axis": 0,
          "time_dim_axis": 1,
          "feature_dim_axis": 3,
          "sparse": False}))
      src_nhwc.output.placeholder = tf_compat.v1.placeholder(shape=(None, None, 16, 16), dtype=tf.float32)
      src_nhwc.output.size_placeholder = {0: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}

    rnd = numpy.random.RandomState(42)
    input_data = rnd.rand(10, 11, 16, 16)
    seq_lens = numpy.array([11] * 10)

    with tf_compat.v1.variable_scope("batch_norm_masked_nchw"):
      batch_norm_1 = BatchNormLayer(
        name="batch_norm_masked_nchw", network=net, masked_time=True,
        sources=[src_nhwc],
        output=BatchNormLayer.get_out_data_from_opts(
          name="batch_norm_masked_nchw",
          sources=[src_nhwc],
          network=net))
      batch_norm_1.post_init(layer_desc={"output": batch_norm_1.output})
    with tf_compat.v1.variable_scope("batch_norm_nonmasked_nchw"):
      batch_norm_2 = BatchNormLayer(
        name="batch_norm_nonmasked_nchw", network=net, masked_time=False,
        sources=[src_nhwc],
        output=BatchNormLayer.get_out_data_from_opts(
          name="batch_norm_nonmasked_nchw",
          sources=[src_nhwc],
          network=net))
      batch_norm_2.post_init(layer_desc={"output": batch_norm_2.output})
    tf_compat.v1.global_variables_initializer().run(session=session)
    out_1, seq_lens_1 = session.run(
      [batch_norm_1.output.placeholder, batch_norm_1.output.size_placeholder[0]],
      feed_dict={
        src_nhwc.output.placeholder: input_data,
        src_nhwc.output.size_placeholder[0]: seq_lens})
    out_2, seq_lens_2 = session.run(
      [batch_norm_2.output.placeholder, batch_norm_2.output.size_placeholder[0]],
      feed_dict={
        src_nhwc.output.placeholder: input_data,
        src_nhwc.output.size_placeholder[0]: seq_lens})
    numpy.testing.assert_array_almost_equal(out_1, out_2)
    print(numpy.sum(out_1 - out_2))


def test_batch_norm_unequal_seq_len():
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData(), train_flag=True)
    with tf_compat.v1.variable_scope("src_nhwc"):
      src_nhwc = InternalLayer(
        name="src_nhwc", network=net,
        output=Data(**{
          "name": "src_nhwc_output",
          "dim": 16,
          "shape": (None, 16, 16),
          "batch_dim_axis": 0,
          "time_dim_axis": 1,
          "feature_dim_axis": 3,
          "sparse": False}))
      src_nhwc.output.placeholder = tf_compat.v1.placeholder(shape=(None, None, 16, 16), dtype=tf.float32)
      src_nhwc.output.size_placeholder = {0: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}

    rnd = numpy.random.RandomState(42)
    input_data = rnd.rand(10, 11, 16, 16).astype('f')
    input_data[2, 5:, :, :] = 0
    input_data_masked = numpy.copy(input_data)
    seq_lens = numpy.array([11, 11, 5, 11, 11, 11, 11, 11, 11, 11], dtype=numpy.float32)
    n1 = 9 * 11 * 16 + 5 * 16
    n2 = 10 * 11 * 16

    with tf_compat.v1.variable_scope("batch_norm_masked_nchw"):
      batch_norm_1 = BatchNormLayer(
        name="batch_norm_masked_nchw", network=net, masked_time=True,
        use_shift=False, use_std=False, epsilon=0.0,
        sources=[src_nhwc],
        output=BatchNormLayer.get_out_data_from_opts(
          name="batch_norm_masked_nchw",
          sources=[src_nhwc],
          network=net))
      batch_norm_1.post_init(layer_desc={"output": batch_norm_1.output})
    with tf_compat.v1.variable_scope("batch_norm_nonmasked_nchw"):
      batch_norm_2 = BatchNormLayer(
        name="batch_norm_nonmasked_nchw", network=net, masked_time=False,
        use_shift=False, use_std=False, epsilon=0,
        sources=[src_nhwc],
        output=BatchNormLayer.get_out_data_from_opts(
          name="batch_norm_nonmasked_nchw",
          sources=[src_nhwc],
          network=net))
      batch_norm_2.post_init(layer_desc={"output": batch_norm_2.output})
    tf_compat.v1.global_variables_initializer().run(session=session)
    out_1, seq_lens_1 = session.run(
      [batch_norm_1.output.placeholder, batch_norm_1.output.size_placeholder[0]],
      feed_dict={
        src_nhwc.output.placeholder: input_data,
        src_nhwc.output.size_placeholder[0]: seq_lens})
    out_2, seq_lens_2 = session.run(
      [batch_norm_2.output.placeholder, batch_norm_2.output.size_placeholder[0]],
      feed_dict={
        src_nhwc.output.placeholder: input_data_masked,
        src_nhwc.output.size_placeholder[0]: seq_lens})

    # Manually calculating batch_norm and compare to the tf output
    data_mean = numpy.mean(input_data, axis=(0, 1, 2), keepdims=True, dtype=numpy.float32)
    data_var = numpy.var(input_data, axis=(0, 1, 2), keepdims=True, dtype=numpy.float32)
    np_bn2 = (input_data - data_mean) * (1.0 / numpy.sqrt(data_var))
    numpy.testing.assert_array_almost_equal(np_bn2, out_2, decimal=5)
    # Manually calculating batch_norm with different seq_lens, having:
    # Mean_1 = n2 / n1 * Mean_2
    # Var_1 = n2 / n1 * (Var_2 + Mean_2 ^ 2 (1 - n2 / n1))
    # bn_1 = (x - Mean_1) * 1 / sqrt(Var_1)
    # Substituting Mean_1 and Var_1:
    np_bn1 = (
      (input_data - n2 / n1 * data_mean) *
      (1.0 / numpy.sqrt(n2 / n1 * (data_var + data_mean ** 2 * (1 - n2 / n1)))))
    # Check with tf output.
    numpy.testing.assert_array_almost_equal(np_bn1, out_1, decimal=5)


def test_BatchNormLayer_CondLayer():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim

  time_dim = SpatialDim('time')
  in_dim = FeatureDim('in', 12)

  config = Config(dict(extern_data={
    'data': {
      'dim_tags': (batch_dim, time_dim, in_dim),
    }
  }))

  net_dict = {
    'length': {
      'class': 'length',
      'from': ['data:data'],
      'axis': batch_dim,
      'out_shape': {}
    },
    'mod': {
      'class': 'eval',
      'from': 'length',
      'eval': 'source(0) % 2',
      'out_shape': {}
    },
    'compare': {
      'class': 'compare',
      'from': 'mod',
      'kind': 'equal',
      'value': 0,
      'out_shape': {}
    },
    'cond': {
      'class': 'cond',
      'from': [],
      'condition': 'compare',
      'true_layer': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'batch_norm': {
            'class': 'subnetwork',
            'from': [],
            'subnetwork': {
              'batch_norm': {
                'class': 'batch_norm',
                'from': 'base:base:data:data',
                'in_dim': in_dim,
                'use_std': True,
                'use_shift': True,
                'param_version': 2,
                'momentum': 0.1,
                'epsilon': 0.001,
                'masked_time': True,
                'out_shape': {batch_dim, time_dim, in_dim}
              },
              'output': {
                'class': 'copy',
                'from': 'batch_norm',
                'out_shape': {batch_dim, time_dim, in_dim}
              }
            },
            'name_scope': ''
          },
          'output': {
            'class': 'copy',
            'from': 'batch_norm',
            'out_shape': {batch_dim, time_dim, in_dim}
          }
        }
      },
      'false_layer': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'output': {
            'class': 'copy',
            'from': 'base:data:data',
            'out_shape': {batch_dim, time_dim, in_dim}
          }
        }
      },
      'out_shape': {batch_dim, time_dim, in_dim},
      'name_scope': ''
    },
    'output': {
      'class': 'copy',
      'from': 'cond',
      'out_shape': {batch_dim, time_dim, in_dim}
    },
  }

  with make_scope() as session:
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict(net_dict)
    net.initialize_params(session)
    out = net.get_default_output_layer().output
    fetches = net.get_fetches_dict()
    fetches["out"] = out.placeholder
    session.run(fetches, feed_dict=make_feed_dict(net.extern_data))


def test_activation_layer_net_construct():
  with make_scope() as session:
    num_inputs = 2
    config = Config()
    config.update({
      "num_outputs": 3,
      "num_inputs": num_inputs,
      "network": {
        "output": {"class": "activation", "activation": "relu", "from": ["data"]}
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    out = network.get_default_output_layer().output.placeholder
    n_batch = 1
    seq_len = 3
    feed = {network.extern_data.get_default_input_data().placeholder:
            numpy.array([[[0, 0], [-1, -1], [2, 2]]], dtype="float32")}
    assert_equal(feed[network.extern_data.get_default_input_data().placeholder].shape, (n_batch, seq_len, num_inputs))
    v = session.run(out, feed_dict=feed)
    assert_equal(v.shape, (n_batch, seq_len, num_inputs))
    assert_equal(v.tolist(), [[[0, 0], [0, 0], [2, 2]]])


def test_activation_layer_abs_for_stft():
  with make_scope() as session:
    num_inputs = 1
    config = Config()
    frame_shift = 1
    frame_size = 3
    fft_size = 3
    config.update({
      "num_outputs": 3,
      "num_inputs": num_inputs,
      "network": {
        "stft": {"class": "multichannel_stft_layer", "from": "data",
                 "frame_shift": frame_shift, "frame_size": frame_size,
                 "window": "hanning", "fft_size": fft_size,
                 "use_rfft": True, "pad_last_frame": False,
                 "nr_of_channels": 1},
        "output": {"class": "activation", "activation": "abs", "from": ["stft"]}
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    out = network.get_default_output_layer().output.placeholder
    n_batch = 1
    seq_len = 6
    feed = {network.extern_data.get_default_input_data().placeholder:
              numpy.array([[[0],[0],[2],[1],[4],[3]]], dtype=numpy.float32)}
    assert_equal(feed[network.extern_data.get_default_input_data().placeholder].shape, (n_batch, seq_len, num_inputs))
    v = session.run(out, feed_dict=feed)
    assert_equal(v.shape, (n_batch, seq_len-(frame_size-1), fft_size//2+1))

    input_stft = tf.signal.stft(numpy.array([[0,0,2,1,4,3]], dtype=numpy.float32), frame_length=frame_size, frame_step=frame_shift, fft_length=fft_size,
                                window_fn=tf.signal.hann_window)
    exp_output = tf.math.abs(input_stft)
    assert_equal(v.tolist(), exp_output.eval().tolist())


def test_activation_layer_net_construct_two_out():
  with make_scope() as session:
    num_inputs = 2
    config = Config()
    config.update({
      "num_outputs": 3,
      "num_inputs": num_inputs,
      "network": {
        "0out": {"class": "linear", "n_out": 1, "activation": "relu", "from": ["data"], "is_output_layer": True},
        "output": {"class": "activation", "activation": "relu", "from": ["data"]}
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    session.run(tf_compat.v1.global_variables_initializer())
    out = network.layers["output"].output.placeholder
    out2 = network.layers["0out"].output.placeholder
    n_batch = 1
    seq_len = 3
    feed = {network.extern_data.get_default_input_data().placeholder:
            numpy.array([[[0, 0], [-1, -1], [2, 2]]], dtype="float32")}
    assert_equal(feed[network.extern_data.get_default_input_data().placeholder].shape, (n_batch, seq_len, num_inputs))
    v, v2 = session.run([out, out2], feed_dict=feed)
    assert_equal(v.shape, (n_batch, seq_len, num_inputs))
    assert_equal(v.tolist(), [[[0, 0], [0, 0], [2, 2]]])


def _test_simple_eval_func(s):
  with make_scope() as session:
    num_inputs = 2
    config = Config()
    config.update({
      "extern_data": {"data": {"dim": num_inputs}},
      "network": {
        "output": {"class": "eval", "eval": "%s(source(0))" % s, "from": "data"}
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    feed = make_feed_dict(network.extern_data)
    out = network.get_default_output_layer().output.placeholder
    session.run(out, feed_dict=feed)


def test_simple_eval_tanh():
  _test_simple_eval_func("tf.tanh")


def test_simple_eval_sigmoid():
  _test_simple_eval_func("tf.sigmoid")


def _test_simple_activation(s):
  with make_scope() as session:
    num_inputs = 2
    config = Config()
    config.update({
      "extern_data": {"data": {"dim": num_inputs}},
      "network": {
        "output": {"class": "activation", "activation": s, "from": "data"}
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    feed = make_feed_dict(network.extern_data)
    out = network.get_default_output_layer().output.placeholder
    session.run(out, feed_dict=feed)


def test_simple_activation_tanh():
  _test_simple_activation("tanh")


def test_simple_activation_log_sigmoid():
  _test_simple_activation("log_sigmoid")


def test_cnn_building_block():
  with make_scope() as session:
    num_inputs = 192
    channel_num = 32
    feature_dim = 6
    filters = 32
    filter_size = (3, 3)
    config = Config()
    config.update({
      "num_inputs": num_inputs,
      "num_outputs": filters,
      "network": {
        "split": {"class": "split_dims", "axis": "f", "dims": (channel_num, feature_dim), "from": ["data"]},
        "swap_axes": {
          "class": "swap_axes", "axis1": "dim:%i" % channel_num, "axis2": "dim:%i" % feature_dim, "from": "split"},
        "c1": {"class": "conv", "n_out": filters, "filter_size": filter_size, "auto_use_channel_first": False,
               "in_spatial_dims": ("T", "dim:6"),
               "strides": (1, 1), "dilation_rate": (1, 1), "padding": "SAME", "activation": None, "with_bias": False,
               "from": "swap_axes"},
        "bn1": {"class": "batch_norm", "from": "c1", "masked_time": True},
        "y1": {"class": "activation", "activation": "relu", "batch_norm": False, "from": "bn1"},
        "c2": {"class": "conv", "n_out": filters, "filter_size": filter_size, "auto_use_channel_first": False,
               "in_spatial_dims": ("T", "dim:6"),
               "strides": (1, 1), "dilation_rate": (1, 1), "padding": "SAME", "activation": None, "with_bias": False,
               "from": "y1"},
        "p": {"class": "combine", "kind": "add", "from": ["c2", "swap_axes"]},
        "bn2": {"class": "batch_norm", "from": "p", "masked_time": True},
        "y2": {"class": "activation", "activation": "relu", "batch_norm": False, "from": "bn2"},

        "out_pool": {
          "class": "reduce", "mode": "avg", "axes": "dim:%i" % feature_dim, "keep_dims": False, "from": "y2"},
        "output": {"class": "copy", "from": ["out_pool"], "is_output_layer": True}
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    swap_layer = network.layers["swap_axes"]
    assert swap_layer.output.dim == channel_num
    session.run(tf_compat.v1.global_variables_initializer())
    out = network.layers["output"].output.placeholder
    n_batch = 5
    seq_len = 10
    seq_lens = numpy.array([10, 10, 10, 10, 10], dtype=numpy.int32)
    feed = {network.extern_data.get_default_input_data().placeholder:
            numpy.random.rand(n_batch, seq_len, num_inputs).astype('f'),
            network.extern_data.get_default_input_data().size_placeholder[0]: seq_lens}
    v = session.run(out, feed_dict=feed)


def test_combine_layer_net_construct():
  with make_scope() as session:
    net_dict = {
      "lstm0_fw": {"class": "rec", "unit": "lstm", "n_out": 5, "direction": 1, "from": "data:data"},
      "lstm0_bw": {"class": "rec", "unit": "lstm", "n_out": 5, "direction": -1, "from": "data:data"},
      "lstm0_avg": {"class": "combine", "kind": "average", "from": ["lstm0_fw", "lstm0_bw"]},
      "output": {"class": "softmax", "loss": "ce", "from": ["lstm0_avg"]}
    }
    config = Config()
    config.update(dict(num_inputs=4, num_outputs=9))
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)


def test_CombineLayer_simple_add():
  with make_scope() as session:
    net_dict = {
      "lin1": {"class": "linear", "activation": "sigmoid", "n_out": 5, "from": "data:data"},
      "lin2": {"class": "linear", "activation": "sigmoid", "n_out": 5, "from": "data:data"},
      "combine": {"class": "combine", "kind": "add", "from": ["lin1", "lin2"]},
      "output": {"class": "softmax", "loss": "ce", "from": "combine"}
    }
    config = Config({"debug_print_layer_output_template": True})
    config.update(dict(num_inputs=4, num_outputs=9))
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    out = network.get_default_output_layer()
    feed_dict = make_feed_dict(network.extern_data.data.values(), same_time=True)
    session.run(tf_compat.v1.global_variables_initializer())
    session.run(out.output.placeholder, feed_dict=feed_dict)


def test_CombineLayer_broadcast():
  with make_scope() as session:
    net_dict = {
      "lin1": {"class": "linear", "activation": "sigmoid", "n_out": 5, "from": "data:data"},
      "lin2": {"class": "linear", "activation": "sigmoid", "n_out": 1, "from": "data:data"},
      "lin2_squeeze": {"class": "squeeze", "from": "lin2", "axis": "f"},
      "combine": {"class": "combine", "kind": "add", "from": ["lin1", "lin2_squeeze"]},
      "output": {"class": "softmax", "loss": "ce", "from": "combine"}
    }
    config = Config({"debug_print_layer_output_template": True})
    config.update(dict(num_inputs=4, num_outputs=9))
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    assert_equal(network.get_layer("combine").output.shape, (None, 5))
    out = network.get_default_output_layer()
    assert out.output.shape == (None, 9)
    feed_dict = make_feed_dict(network.extern_data.data.values(), same_time=True)
    session.run(tf_compat.v1.global_variables_initializer())
    session.run(out.output.placeholder, feed_dict=feed_dict)


def test_CombineLayer_broadcast_multiple():
  with make_scope() as session:
    net_dict = {
      "p1": {"class": "variable", "shape": (5, 5, 3), "add_batch_axis": False},
      "p2": {"class": "variable", "shape": (5,), "add_batch_axis": False},
      "combine": {"class": "combine", "kind": "add", "from": ["p1", "p2"]},
      "output": {"class": "softmax", "loss": "ce", "from": "combine"}
    }
    config = Config({"debug_print_layer_output_template": True})
    config.update(dict(num_inputs=4, num_outputs=9))
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    assert_equal(network.get_layer("combine").output.batch_shape, (5, 5, 3))
    out = network.get_default_output_layer()
    assert out.output.batch_shape == (5, 5, 9) and not out.output.have_batch_axis()
    feed_dict = make_feed_dict(network.extern_data.data.values(), same_time=True)
    session.run(tf_compat.v1.global_variables_initializer())
    out_v = session.run(out.output.placeholder, feed_dict=feed_dict)
    assert out_v.shape == out.output.batch_shape


def test_CombineLayer_broadcast_same_dim_diff_tag():
  from returnn.tf.util.data import batch_dim, FeatureDim, SpatialDim
  time_dim = SpatialDim("time")
  input_dim = FeatureDim("input", 3)
  feat_dim = FeatureDim("feat", 3)

  template_net = TFNetwork(config=Config(), extern_data=ExternData())
  out = CombineLayer.get_out_data_from_opts(
    network=template_net, name="template_combine", kind="add",
    sources=[
      InternalLayer(name="a", network=template_net, output=Data("a", dim_tags=[batch_dim, time_dim, input_dim])),
      InternalLayer(name="b", network=template_net, output=Data("b", dim_tags=[feat_dim])),
    ],
    allow_broadcast_all_sources=True)
  assert out.dim_tags == (batch_dim, time_dim, input_dim, feat_dim)

  with make_scope() as session:
    net_dict = {
      "p": {"class": "variable", "shape": [feat_dim], "add_batch_axis": False},
      "output": {"class": "combine", "kind": "add", "from": ["data", "p"], "allow_broadcast_all_sources": True},
    }
    config = Config({
      "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, input_dim]}}})
    network = TFNetwork(config=config)
    network.construct_from_dict(net_dict)
    out = network.get_default_output_layer()
    assert out.output.dim_tags == (batch_dim, time_dim, input_dim, feat_dim)
    feed_dict = make_feed_dict(network.extern_data, n_batch=2, n_time=5)
    session.run(tf_compat.v1.global_variables_initializer())
    out_v = session.run(out.output.placeholder, feed_dict=feed_dict)
    assert out_v.shape == (2, 5, 3, 3)


def test_CombineLayer_match_unknown():
  with make_scope() as session:
    dat1 = Data(name="undefined", shape=(None, 3))
    assert dat1.dim_tags[1].undefined
    # Create placeholders to have this dyn size clearly defined.
    dat2 = Data(name="defined", shape=(None, 3), auto_create_placeholders=True)
    net = TFNetwork(extern_data=ExternData())
    layer1 = InternalLayer(name="layer1_undefined", network=net, output=dat1)
    layer2 = InternalLayer(name="layer2_defined", network=net, output=dat2)
    out = CombineLayer.get_out_data_from_opts(name="combine", network=net, sources=[layer1, layer2])
    assert out.dim_tags[:2] == dat2.dim_tags[:2] and out.batch_shape == dat2.batch_shape


def test_CombineLayer_match_unknown_derived():
  with make_scope() as session:
    dat1 = Data(name="undefined", shape=(None, 3))
    assert dat1.dim_tags[1].undefined
    dat1_derived_dim_tags = list(dat1.dim_tags)
    dat1_derived_dim_tags[1] = Dim(
      kind=Dim.Types.Spatial, description="undefined_derived_dim", derived_from_tag=dat1.dim_tags[1])
    dat1_derived = Data(name="undefined_derived", dim_tags=dat1_derived_dim_tags)
    assert dat1_derived.dim_tags[1].undefined
    # Create placeholders to have this dyn size clearly defined.
    dat2 = Data(name="defined", shape=(None, 3), auto_create_placeholders=True)
    net = TFNetwork(extern_data=ExternData())
    layer1 = InternalLayer(name="layer1_undefined_derived", network=net, output=dat1_derived)
    layer2 = InternalLayer(name="layer2_defined", network=net, output=dat2)
    out = CombineLayer.get_out_data_from_opts(name="combine", network=net, sources=[layer1, layer2])
    assert out.dim_tags[:2] == dat2.dim_tags[:2] and out.batch_shape == dat2.batch_shape


def test_CombineLayer_match_unknown_batch_feature_major():
  with make_scope() as session:
    dat1 = Data(name="undefined", shape=(None, 1))
    assert dat1.dim_tags[1].undefined
    # Create placeholders to have this dyn size clearly defined.
    dat2 = Data(name="defined", shape=(None, 1), auto_create_placeholders=True)
    dat2_bf_major = dat2.copy_as_batch_feature_major()
    net = TFNetwork(extern_data=ExternData())
    layer1 = InternalLayer(name="layer1_undefined", network=net, output=dat1)
    layer2 = InternalLayer(name="layer2_defined", network=net, output=dat2_bf_major)
    out = CombineLayer.get_out_data_from_opts(name="combine", network=net, sources=[layer1, layer2])
    assert out.dim_tags[:2] == dat2.dim_tags[:2] and out.batch_shape == dat2.batch_shape


def test_CombineLayer_match_unknown_batch_feature_major_with_out_type():
  with make_scope() as session:
    dat1 = Data(name="undefined", shape=(None, 1))
    assert dat1.dim_tags[1].undefined
    # Create placeholders to have this dyn size clearly defined.
    dat2 = Data(name="defined", shape=(None, 1), auto_create_placeholders=True)
    dat2_bf_major = dat2.copy_as_batch_feature_major()
    net = TFNetwork(extern_data=ExternData())
    layer1 = InternalLayer(name="layer1_undefined", network=net, output=dat1)
    layer2 = InternalLayer(name="layer2_defined", network=net, output=dat2_bf_major)
    out = CombineLayer.get_out_data_from_opts(
      name="combine", network=net, sources=[layer1, layer2],
      out_type={'dim': 1, 'shape': (None, 1)})
    assert out.dim_tags[:2] == dat2.dim_tags[:2] and out.batch_shape == dat2.batch_shape


def test_CombineLayer_different_batch_axis():
  # ["base:enc_ctx", "weight_feedback", "s_transformed"]
  # base:enc_ctx: Data(name='enc_ctx_output', shape=(None, 14), batch_dim_axis=1)
  # weight_feedback: Data(name='weight_feedback_output', shape=(None, 14), batch_dim_axis=1)
  # s_transformed: Data(name='s_transformed_output', shape=(14,), time_dim_axis=None)
  # out: Data(name='energy_in_output', shape=(None, 14), beam_size=3)
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  feature_dim = FeatureDim("feature", 5)
  time_dim = SpatialDim("time")
  with make_scope() as session:
    config = Config({
      "extern_data": {
        "enc_ctx": {"dim_tags": (batch_dim, time_dim, feature_dim), "available_for_inference": True},
        "weight_feedback": {"dim_tags": (time_dim, batch_dim, feature_dim), "available_for_inference": True},
        "s_transformed": {"dim_tags": (batch_dim, feature_dim), "available_for_inference": True},
      }
    })
    net = TFNetwork(config=config, train_flag=True)
    l1 = net.get_layer("data:enc_ctx")
    l2 = net.get_layer("data:weight_feedback")
    l3 = net.get_layer("data:s_transformed")
    out = net.add_layer(name="energy_in", layer_class=CombineLayer, kind="add", sources=[l1, l2, l3])
    print("out:", out)
    assert out.output.dim_tags == (batch_dim, time_dim, feature_dim)
    session.run(out.output.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_CombineLayer_two_time_dims():
  with make_scope() as session:
    n_dim = 5
    n_batch = 3
    n_time1 = 7
    n_time2 = 11
    rnd = numpy.random.RandomState(42)
    net_dict = {
      "output": {
        "class": "combine", "kind": "add",
        "from": ["data:in0", "data:in1", "data:in2"]}
    }
    config = Config({"debug_print_layer_output_template": True})
    extern_data = ExternData()
    in0 = Data(
      name="in0", shape=(None, None, n_dim), batch_dim_axis=1, auto_create_placeholders=True)
    in1 = Data(
      # same time as first in in0
      name="in1", shape=(None, n_dim), auto_create_placeholders=True)
    in2 = Data(
      # same time as in second in in0
      name="in2", shape=(None, n_dim), batch_dim_axis=1, auto_create_placeholders=True)
    extern_data.register_data(in0)
    extern_data.register_data(in1)
    extern_data.register_data(in2)
    in1.get_size_dim_tag(0).declare_same_as(in0.get_size_dim_tag(0))
    in2.get_size_dim_tag(0).declare_same_as(in0.get_size_dim_tag(1))
    print("ExternData all dimension tags (allow_same_feature_dim=True):")
    pprint(extern_data.get_all_dimension_tags(allow_same_feature_dim=True))
    network = TFNetwork(config=config, extern_data=extern_data, train_flag=True)
    network.construct_from_dict(net_dict)
    output = network.get_default_output_layer().output
    assert output.shape == (None, None, n_dim) and set(output.size_placeholder.keys()) == {0, 1}
    assert output.batch_dim_axis == 1 and output.time_dim_axis == 0
    time1_np = numpy.array([n_time1, n_time1 - 3, n_time1 - 2])
    assert min(time1_np) > 0 and max(time1_np) == n_time1 and len(time1_np) == n_batch
    time2_np = numpy.array([n_time2, n_time2 - 2, n_time2 - 5])
    assert min(time2_np) > 0 and max(time2_np) == n_time2 and len(time2_np) == n_batch
    in0_np = rnd.normal(size=(n_time1, n_batch, n_time2, n_dim)).astype("float32")
    in1_np = rnd.normal(size=(n_batch, n_time1, n_dim)).astype("float32")
    in2_np = rnd.normal(size=(n_time2, n_batch, n_dim)).astype("float32")
    out_np, out_sizes_np = session.run(
      fetches=(output.placeholder, output.size_placeholder.as_dict()),
      feed_dict={
        in0.placeholder: in0_np, in0.size_placeholder[0]: time1_np, in0.size_placeholder[1]: time2_np,
        in1.placeholder: in1_np, in1.size_placeholder[0]: time1_np,
        in2.placeholder: in2_np, in2.size_placeholder[0]: time2_np})
    assert isinstance(out_np, numpy.ndarray)
    assert isinstance(out_sizes_np, dict) and set(out_sizes_np.keys()) == {0, 1}
    out_time0_np, out_time1_np = out_sizes_np[0], out_sizes_np[1]
    assert isinstance(out_time0_np, numpy.ndarray) and isinstance(out_time1_np, numpy.ndarray)
    assert out_np.shape == (n_time1, n_batch, n_time2, n_dim)


def test_CombineLayer_two_time_dims_first_not_most_generic():
  with make_scope() as session:
    n_dim = 5
    n_batch = 3
    n_time1 = 7
    n_time2 = 11
    rnd = numpy.random.RandomState(42)
    net_dict = {
      "output": {
        "class": "combine", "kind": "add",
        "from": ["data:in1", "data:in0", "data:in2"]}
    }
    config = Config({"debug_print_layer_output_template": True})
    extern_data = ExternData()
    in0 = Data(
      name="in0", shape=(None, None, n_dim), batch_dim_axis=1, auto_create_placeholders=True)
    in1 = Data(
      # same time as first in in0
      name="in1", shape=(None, n_dim), auto_create_placeholders=True)
    in2 = Data(
      # same time as in second in in0
      name="in2", shape=(None, n_dim), batch_dim_axis=1, auto_create_placeholders=True)
    extern_data.register_data(in0)
    extern_data.register_data(in1)
    extern_data.register_data(in2)
    in1.get_size_dim_tag(0).declare_same_as(in0.get_size_dim_tag(0))
    in2.get_size_dim_tag(0).declare_same_as(in0.get_size_dim_tag(1))
    print("ExternData all dimension tags (allow_same_feature_dim=True):")
    pprint(extern_data.get_all_dimension_tags(allow_same_feature_dim=True))
    network = TFNetwork(config=config, extern_data=extern_data, train_flag=True)
    network.construct_from_dict(net_dict)
    output = network.get_default_output_layer().output
    assert output.shape == (None, None, n_dim) and set(output.size_placeholder.keys()) == {0, 1}
    assert output.batch_dim_axis == 1 and output.time_dim_axis == 0
    time1_np = numpy.array([n_time1, n_time1 - 3, n_time1 - 2])
    assert min(time1_np) > 0 and max(time1_np) == n_time1 and len(time1_np) == n_batch
    time2_np = numpy.array([n_time2, n_time2 - 2, n_time2 - 5])
    assert min(time2_np) > 0 and max(time2_np) == n_time2 and len(time2_np) == n_batch
    in0_np = rnd.normal(size=(n_time1, n_batch, n_time2, n_dim)).astype("float32")
    in1_np = rnd.normal(size=(n_batch, n_time1, n_dim)).astype("float32")
    in2_np = rnd.normal(size=(n_time2, n_batch, n_dim)).astype("float32")
    out_np, out_sizes_np = session.run(
      fetches=(output.placeholder, output.size_placeholder.as_dict()),
      feed_dict={
        in0.placeholder: in0_np, in0.size_placeholder[0]: time1_np, in0.size_placeholder[1]: time2_np,
        in1.placeholder: in1_np, in1.size_placeholder[0]: time1_np,
        in2.placeholder: in2_np, in2.size_placeholder[0]: time2_np})
    assert isinstance(out_np, numpy.ndarray)
    assert isinstance(out_sizes_np, dict) and set(out_sizes_np.keys()) == {0, 1}
    out_time0_np, out_time1_np = out_sizes_np[0], out_sizes_np[1]
    assert isinstance(out_time0_np, numpy.ndarray) and isinstance(out_time1_np, numpy.ndarray)
    assert out_np.shape == (n_time1, n_batch, n_time2, n_dim)


def test_CombineLayer_two_time_dims_first_not_most_generic_with_n_out():
  with make_scope() as session:
    n_dim = 5
    n_batch = 3
    n_time1 = 7
    n_time2 = 11
    rnd = numpy.random.RandomState(42)
    net_dict = {
      "output": {
        "class": "combine", "kind": "add", "n_out": n_dim,
        "from": ["data:in1", "data:in0", "data:in2"]}
    }
    config = Config({"debug_print_layer_output_template": True})
    extern_data = ExternData()
    in0 = Data(
      name="in0", shape=(None, None, n_dim), batch_dim_axis=1, auto_create_placeholders=True)
    in1 = Data(
      # same time as first in in0
      name="in1", shape=(None, n_dim), auto_create_placeholders=True)
    in2 = Data(
      # same time as in second in in0
      name="in2", shape=(None, n_dim), batch_dim_axis=1, auto_create_placeholders=True)
    extern_data.register_data(in0)
    extern_data.register_data(in1)
    extern_data.register_data(in2)
    in1.get_size_dim_tag(0).declare_same_as(in0.get_size_dim_tag(0))
    in2.get_size_dim_tag(0).declare_same_as(in0.get_size_dim_tag(1))
    print("ExternData all dimension tags (allow_same_feature_dim=True):")
    pprint(extern_data.get_all_dimension_tags(allow_same_feature_dim=True))
    network = TFNetwork(config=config, extern_data=extern_data, train_flag=True)
    network.construct_from_dict(net_dict)
    output = network.get_default_output_layer().output
    assert output.shape == (None, None, n_dim) and set(output.size_placeholder.keys()) == {0, 1}
    assert output.batch_dim_axis == 1 and output.time_dim_axis == 0
    time1_np = numpy.array([n_time1, n_time1 - 3, n_time1 - 2])
    assert min(time1_np) > 0 and max(time1_np) == n_time1 and len(time1_np) == n_batch
    time2_np = numpy.array([n_time2, n_time2 - 2, n_time2 - 5])
    assert min(time2_np) > 0 and max(time2_np) == n_time2 and len(time2_np) == n_batch
    in0_np = rnd.normal(size=(n_time1, n_batch, n_time2, n_dim)).astype("float32")
    in1_np = rnd.normal(size=(n_batch, n_time1, n_dim)).astype("float32")
    in2_np = rnd.normal(size=(n_time2, n_batch, n_dim)).astype("float32")
    out_np, out_sizes_np = session.run(
      fetches=(output.placeholder, output.size_placeholder.as_dict()),
      feed_dict={
        in0.placeholder: in0_np, in0.size_placeholder[0]: time1_np, in0.size_placeholder[1]: time2_np,
        in1.placeholder: in1_np, in1.size_placeholder[0]: time1_np,
        in2.placeholder: in2_np, in2.size_placeholder[0]: time2_np})
    assert isinstance(out_np, numpy.ndarray)
    assert isinstance(out_sizes_np, dict) and set(out_sizes_np.keys()) == {0, 1}
    out_time0_np, out_time1_np = out_sizes_np[0], out_sizes_np[1]
    assert isinstance(out_time0_np, numpy.ndarray) and isinstance(out_time1_np, numpy.ndarray)
    assert out_np.shape == (n_time1, n_batch, n_time2, n_dim)


def test_CombineLayer_two_time_dims_allow_broadcast_all_sources():
  from returnn.tf.util.data import batch_dim
  with make_scope() as session:
    n_dim = 5
    n_batch = 3
    n_time1 = 7
    n_time2 = 11
    time1_dim = SpatialDim("time1")
    time2_dim = SpatialDim("time2")
    feat_dim = FeatureDim("feature", dimension=n_dim)
    rnd = numpy.random.RandomState(42)
    config = Config({"debug_print_layer_output_template": True})
    extern_data = ExternData()
    in1 = Data(name="in1", dim_tags=[batch_dim, time1_dim, feat_dim], auto_create_placeholders=True)
    in2 = Data(name="in2", dim_tags=[batch_dim, time2_dim, feat_dim], auto_create_placeholders=True)
    extern_data.register_data(in1)
    extern_data.register_data(in2)
    print("ExternData all dimension tags (allow_same_feature_dim=True):")
    pprint(extern_data.get_all_dimension_tags(allow_same_feature_dim=True))
    network = TFNetwork(config=config, extern_data=extern_data, train_flag=True)
    try:
      network.construct_from_dict({
        "output": {"class": "combine", "kind": "add", "from": ["data:in1", "data:in2"]}})
    except Exception as exc:
      # https://github.com/rwth-i6/returnn/issues/691
      print("Expected exception:", exc)
      assert "require broadcasting" in str(exc)
    else:
      raise Exception(
        "Expect allow_broadcast_all_sources exception, but layer constructed: %s" % network.get_default_output_layer())
    network.construct_from_dict({
      "output": {
        "class": "combine", "kind": "add", "from": ["data:in1", "data:in2"],
        "out_shape": {batch_dim, time1_dim, time2_dim, feat_dim}}})
    output = network.get_default_output_layer().output
    assert output.shape == (None, None, n_dim) and set(output.size_placeholder.keys()) == {0, 1}
    time1_np = numpy.array([n_time1, n_time1 - 3, n_time1 - 2])
    assert min(time1_np) > 0 and max(time1_np) == n_time1 and len(time1_np) == n_batch
    time2_np = numpy.array([n_time2, n_time2 - 2, n_time2 - 5])
    assert min(time2_np) > 0 and max(time2_np) == n_time2 and len(time2_np) == n_batch
    in1_np = rnd.normal(size=(n_batch, n_time1, n_dim)).astype("float32")
    in2_np = rnd.normal(size=(n_batch, n_time2, n_dim)).astype("float32")
    out_np, out_sizes_np = session.run(
      fetches=(output.placeholder, output.size_placeholder.as_dict()),
      feed_dict={
        in1.placeholder: in1_np, in1.size_placeholder[0]: time1_np,
        in2.placeholder: in2_np, in2.size_placeholder[0]: time2_np})
    assert isinstance(out_np, numpy.ndarray)
    assert isinstance(out_sizes_np, dict) and set(out_sizes_np.keys()) == {0, 1}
    out_time0_np, out_time1_np = out_sizes_np[0], out_sizes_np[1]
    assert isinstance(out_time0_np, numpy.ndarray) and isinstance(out_time1_np, numpy.ndarray)
    assert out_np.shape == (n_batch, n_time1, n_time2, n_dim)


def test_CombineLayer_time_broadcast():
  with make_scope() as session:
    n_batch, n_time, n_features = 3, 7, 5
    net_dict = {
      "output": {"class": "combine", "kind": "add", "from": ["data:in1", "data:in2"]},
    }
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "in1": {"shape": (n_features,), "batch_dim_axis": None, "time_dim_axis": None, "feature_dim_axis": 0},
        "in2": {"shape": (n_features, None), "batch_dim_axis": 0, "time_dim_axis": 2}
      }
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    out = network.get_default_output_layer()
    assert_equal(out.output.batch_shape, (None, n_features, None))
    feed_dict = make_feed_dict(network.extern_data, n_batch=n_batch, n_time=n_time)
    session.run(tf_compat.v1.global_variables_initializer())
    out_v = session.run(out.output.placeholder, feed_dict=feed_dict)
    assert out_v.shape == (n_batch, n_features, n_time)


def test_CombineLayer_time_broadcast_swapped():
  with make_scope() as session:
    n_batch, n_time, n_features = 3, 7, 5
    net_dict = {
      "output": {"class": "combine", "kind": "add", "from": ["data:in1", "data:in2"]},
    }
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "in1": {"shape": (n_features, None), "batch_dim_axis": 0, "time_dim_axis": 2},
        "in2": {"shape": (n_features,), "batch_dim_axis": None, "time_dim_axis": None, "feature_dim_axis": 0},
      }
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    out = network.get_default_output_layer()
    assert_equal(out.output.batch_shape, (None, n_features, None))
    feed_dict = make_feed_dict(network.extern_data, n_batch=n_batch, n_time=n_time)
    session.run(tf_compat.v1.global_variables_initializer())
    out_v = session.run(out.output.placeholder, feed_dict=feed_dict)
    assert out_v.shape == (n_batch, n_features, n_time)


def test_CombineLayer_RangeFromLengthLayer():
  from returnn.tf.util.basic import py_print
  from returnn.tf.util.data import batch_dim, Dim, ImplicitDynSizeDim

  def _eval_seq_lens(source, **_kwargs):
    # Get some random varying seq lens.
    res = tf.cast(6.3 * tf.maximum(source(0), 0.), tf.int32) + 1
    res = py_print(res, ["seq lens", res])
    return res

  time_dim = SpatialDim("T")
  new_time_dim = SpatialDim("T_new")
  feat_dim = FeatureDim("F", dimension=13)
  net_dict = {
    "data_red1": {
      "class": "reduce", "from": "data", "axis": "T", "mode": "mean", "out_shape": {batch_dim, feat_dim}},
    "data_red2": {
      "class": "reduce", "from": "data_red1", "axis": "F", "mode": "sum", "out_shape": {batch_dim}},
    "seq_lens": {
      "class": "eval", "from": "data_red2", "eval": _eval_seq_lens,
      "out_type": {"dtype": "int32"}, "out_shape": {batch_dim}},
    "range": {
      "class": "range_from_length", "from": "seq_lens", "out_spatial_dim": new_time_dim,
      "out_shape": {new_time_dim, ImplicitDynSizeDim(batch_dim)}},
    "combine": {
      "class": "eval", "from": ["data_red1", "range"],
      "eval": "source(0) + 0.1 * tf.cast(source(1), tf.float32)",
      "out_shape": {batch_dim, new_time_dim, feat_dim}},
    "output": {"class": "copy", "from": "combine"},
  }

  config = Config({
    "debug_print_layer_output_template": True,
    "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}},
  })

  with make_scope() as session:
    net = TFNetwork(config=config)
    in_data = net.extern_data.get_default_input_data()
    net.construct_from_dict(net_dict)
    out_data = net.get_default_output_layer().output
    seq_lens_data = net.get_layer("seq_lens").output
    assert out_data.batch_ndim == in_data.batch_ndim == 3
    assert out_data.dim_tags[0] == in_data.dim_tags[0] and out_data.dim_tags[0].is_batch_dim()
    out_time = out_data.get_time_dim_tag()
    assert out_time.dyn_size_ext.placeholder is seq_lens_data.placeholder
    assert out_time not in in_data.dim_tags
    feed_dict = make_feed_dict(net.extern_data)
    session.run((out_data.placeholder, out_data.get_sequence_lengths()), feed_dict=feed_dict)


def test_CompareLayer_allow_broadcast_all_sources():
  from returnn.tf.util.data import batch_dim, Dim
  time_tag = Dim(Dim.Types.Spatial, description="time")
  with make_scope():
    n_out = 5
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out, "same_dim_tags_as": {"t": time_tag}}
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "range0": {"class": "range_in_axis", "from": "data", "axis": "b"},
      "range1": {"class": "range_in_axis", "from": "data", "axis": "t"},
      "compare": {
        "class": "compare", "from": ["range0", "range1"], "kind": "equal", "is_output_layer": True,
        "out_shape": {batch_dim, time_tag}}
    })


def test_RangeFromLength_over_batch():
  # https://github.com/rwth-i6/pytorch-to-returnn/issues/100
  net_dict = {
    "batch_len": {'class': 'length', 'axis': 'B', 'from': 'data'},
    "output": {'class': 'range_from_length', 'from': 'batch_len'},
  }
  with make_scope() as session:
    config = Config({"extern_data": {"data": {"dim": 3}}})
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    in_ = net.extern_data.get_default_input_data()
    out = net.get_default_output_layer().output
    assert out.dim_tags[0].is_batch_dim()
    in_v, out_v = session.run((in_.placeholder, out.placeholder), feed_dict=make_feed_dict(net.extern_data))
    n_batch, _, _ = in_v.shape
    assert out_v.shape == (n_batch,)
    assert list(out_v) == list(range(n_batch))


def test_RangeInAxisLayer():
  net_dict = {
    "output": {'class': 'range_in_axis', 'from': 'data', 'axis': 'T'},
  }
  with make_scope() as session:
    config = Config({"extern_data": {"data": {"dim": 3}}})
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    in_ = net.extern_data.get_default_input_data()
    out = net.get_default_output_layer().output
    assert out.get_time_dim_tag() == in_.get_time_dim_tag()
    in_v, out_v = session.run((in_.placeholder, out.placeholder), feed_dict=make_feed_dict(net.extern_data))
    n_batch, n_time, n_feat = in_v.shape
    assert out_v.shape == (n_time,)
    assert list(out_v) == list(range(n_time))


def test_RangeInAxisLayer_generic_dim():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  out_time_dim = time_dim + 2
  feat_dim = FeatureDim("feat", 3)
  config = Config({"extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}}})
  net_dict = {
    "output": {'class': 'range_in_axis', 'from': 'data', 'axis': out_time_dim},
  }
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    in_ = net.extern_data.get_default_input_data()
    out = net.get_default_output_layer().output
    assert in_.get_time_dim_tag() == time_dim != out_time_dim == out.get_time_dim_tag()
    in_v, out_v = session.run((in_.placeholder, out.placeholder), feed_dict=make_feed_dict(net.extern_data))
    n_batch, n_time, n_feat = in_v.shape
    assert out_v.shape == (n_time + 2,)
    assert list(out_v) == list(range(n_time + 2))


def test_SwitchLayer_sanity_check():
  """
  https://github.com/rwth-i6/returnn/issues/800
  """
  from returnn.tf.util.data import Dim
  with make_scope():
    n_out = 5
    time_tag = Dim(Dim.Types.Spatial, "time")
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out, "same_dim_tags_as": {"t": time_tag}},
        "classes": {"dim": n_out, "sparse": True, "same_dim_tags_as": {"t": time_tag}}
      }})
    net = TFNetwork(config=config, search_flag=True)
    net.construct_from_dict({
      "data_int": {"class": "cast", "from": "data", "dtype": "int32"},
      "output": {"class": "rec", "from": "data", "unit": {
        "prev_out": {
          "class": "reinterpret_data", "from": "prev:output", "set_sparse": False},
        "cond": {
          "class": "compare", "from": ["prev_out", "base:data_int"], "kind": "less_equal"},
        "switch": {"class": "switch", "condition": "cond", "true_from": 0.0, "false_from": 0.0},
        "switch_red": {"class": "reduce", "from": "switch", "axis": "t", "mode": "max"},
        "output_prob": {"class": "softmax", "from": "switch_red", "target": "classes", "loss": "ce"},
        "output": {
          "class": "choice", "from": "output_prob", "beam_size": 3, "target": "classes",
          "input_type": "prob", "initial_output": 0}
      }}
    })


def test_dot_layer_shuffled_remaining_dims_static():
  with make_scope() as session:
    import numpy as np
    net_dict = {
      "a": {"class": "split_dims", "axis": "F", "dims": (2, 3, 5), "from": "data:data"},
      "b": {"class": "transpose", "from": ["a"], "perm": {"dim:2": "dim:3", "dim:3": "dim:2"}},
      "dot": {
        "class": "dot", "from": ["a", "b"],
        "red1": "dim:5", "red2": "dim:5", "var1": None, "var2": None,
        "debug": True},
      "output": {"class": "merge_dims", "axes": ["dim:2", "dim:3"], "from": "dot"}
    }
    config = Config()
    config.update({
      "extern_data": {"data": {"shape": (30,)}},
      "network": net_dict
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    out = network.get_default_output_layer(must_exist=True)
    input_data = np.ones(shape=(17, 30))
    feed_dict = {network.extern_data.data["data"].placeholder: input_data}

    # just check that it runs
    session.run(out.output.placeholder, feed_dict)


def test_dot_layer_shuffled_remaining_dims_dynamic():
  with make_scope() as session:
    import numpy as np
    batch_size = 8
    time_size = 20
    feat_size = 10
    net_dict = {
      "a": {"class": "copy", "from": "data"},
      "b": {"class": "transpose", "from": ["a"], "perm": {"B": "T", "T": "B"}},
      "dot": {
        "class": "dot", "from": ["a", "b"],
        "red1": "F", "red2": "F", "var1": None, "var2": None,
        "debug": True},
      "output": {"class": "copy", "from": ["dot"]}
    }
    config = Config()
    config.update({
      "num_outputs": 1,
      "num_inputs": feat_size,
      "network": net_dict
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    out = network.get_default_output_layer(must_exist=True)
    input_data = np.ones(shape=(batch_size, time_size, feat_size))
    feed_dict = {network.layers['data'].output.placeholder: input_data}

    # just check that it runs
    session.run(out.output.placeholder, feed_dict)


def test_dropout_layer_net_construct():
  with make_scope() as session:
    net_dict = {
      "drop": {"class": "dropout", "dropout": 0.3, "dropout_noise_shape": {"*": None}, "from": "data:data"},
      "output": {"class": "softmax", "loss": "ce", "from": ["drop"]}
    }
    config = Config({"num_inputs": 4, "num_outputs": 9, "debug_print_layer_output_template": True})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)


def test_DropoutLayer_axis():
  with make_scope() as session:
    net_dict = {
      "drop": {"class": "dropout", "dropout": 0.3, "dropout_axis": ["T", "F"], "from": "data:data"},
      "output": {"class": "softmax", "loss": "ce", "from": "drop"}
    }
    config = Config({"num_inputs": 4, "num_outputs": 9})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    network.initialize_params(session)
    session.run(network.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(network.extern_data))


def test_ScaledGradientLayer_tensor():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim

  time_dim = SpatialDim('time')
  input_dim = FeatureDim('input', 3)
  out_dim = FeatureDim('out', 5)

  config = Config(dict(extern_data={
    'data': {'dim_tags': (batch_dim, time_dim, input_dim)},
    'classes': {'dim_tags': (batch_dim, time_dim), 'sparse_dim': out_dim},
  }))

  net_dict = {
    'weight_raw': {
      'class': 'variable',
      'shape': [input_dim, out_dim],
      'param_name': 'param',
    },
    'scales': {
      'class': 'random', 'from': [], 'shape': [out_dim],
      'distribution': 'uniform', 'minval': 0.1, 'maxval': 1.5,
    },
    'weight': {
      'class': 'scaled_grad', 'from': 'weight_raw', 'scale': 'scales',
    },
    'bias_raw': {
      'class': 'variable',
      'shape': [out_dim],
      'param_name': 'param',
    },
    'bias': {
      'class': 'scaled_grad', 'from': 'bias_raw', 'scale': 'scales',
    },
    'dot': {
      'class': 'dot',
      'from': ['data:data', 'weight'],
      'reduce': input_dim,
      'out_shape': {batch_dim, time_dim, out_dim}
    },
    'add': {
      'class': 'combine',
      'from': ['dot', 'bias'],
      'kind': 'add',
      'out_shape': {batch_dim, time_dim, out_dim}
    },
    'ce': {
      'class': 'sparse_softmax_cross_entropy_with_logits',
      'logits': 'add',
      'targets': 'data:classes',
      'axis': out_dim,
      'loss': 'as_is',
      'out_shape': {batch_dim, time_dim}
    },
  }

  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    optimizer = tf_compat.v1.train.AdamOptimizer(0.1)
    loss = network.get_total_loss()
    optim_op = optimizer.minimize(network.get_objective())
    session.run(tf_compat.v1.global_variables_initializer())
    feed_dict = make_feed_dict(network.extern_data)
    last_loss_v = None
    for step in range(10):
      loss_v, _ = session.run((loss, optim_op), feed_dict=feed_dict)
      print("Step %i, loss %f" % (step, loss_v))
      if last_loss_v is not None:
        assert loss_v < last_loss_v
      last_loss_v = loss_v


def test_subnetwork_layer_net_construct():
  with make_scope() as session:
    net_dict = {
      "ff0": {"class": "forward", "activation": "tanh", "n_out": 3, "from": "data:data"},
      "sub": {"class": "subnetwork", "from": "ff0", "subnetwork": {
        "ff1": {"class": "forward", "activation": "relu", "n_out": 2, "from": "data"},  # unused
        "output": {"class": "forward", "activation": "relu", "n_out": 2, "from": "data"}
      }},
      "output": {"class": "softmax", "loss": "ce", "from": "sub"}
    }
    config = Config()
    config.update(dict(num_inputs=4, num_outputs=3))
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    assert_equal(network.layers["sub"].output.dim, 2)
    sub_layer = network.layers["sub"]
    assert isinstance(sub_layer, SubnetworkLayer)
    sub_layer_deps = sub_layer.get_dep_layers()
    assert sub_layer_deps, "%r no deps" % sub_layer
    all_deps = set()
    queue = [sub_layer]
    while queue:
      layer = queue.pop(0)
      if layer in all_deps:
        continue
      all_deps.add(layer)
      for dep in layer.get_dep_layers():
        if dep not in all_deps:
          queue.append(dep)
    assert network.layers["ff0"] in all_deps


def test_subnet_loss():
  with make_scope() as session:
    config = Config({
      "extern_data": {"data": {"dim": 1}},
      "debug_print_layer_output_template": True
    })
    net_dict = {
      "sub": {"class": "subnetwork", "from": [], "subnetwork": {
        "var": {"class": "variable", "shape": [1]},
        "loss": {"class": "copy", "from": "var", "loss": "as_is"},
        "output": {"class": "copy", "from": "var"}  # no dep on "loss"
      }},
      # Output dep on "sub" to trigger subnet creation.
      # In theory, it would be nice if the loss is also constructed without that,
      # but this doesn't work currently as "sub" is never constructed by the current heuristics.
      "output": {"class": "copy", "from": "sub"}
    }
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    losses_dict, total_loss, total_constraints = network.get_losses_initialized(with_total=True)
    print("losses:")
    pprint(losses_dict)
    assert len(losses_dict) == 1 and set(losses_dict.keys()) == {"sub/loss"}


def test_subnet2_loss():
  with make_scope() as session:
    config = Config({
      "extern_data": {"data": {"dim": 1}},
      "debug_print_layer_output_template": True
    })
    net_dict = {
      "sub": {"class": "subnetwork", "from": [], "subnetwork": {
        "var": {"class": "variable", "shape": [1]},
        "loss": {"class": "copy", "from": "var", "loss": "as_is"},
        "output": {"class": "copy", "from": "var"}  # no dep on "loss"
      }},
      # Output dep on "sub" to trigger subnet creation.
      # In theory, it would be nice if the loss is also constructed without that,
      # but this doesn't work currently as "sub" is never constructed by the current heuristics.
      # Specifically depend on "sub/output", because in that case,
      # the SubnetworkLayer itself might not be created with the new subnet logic,
      # which is sth we want to test.
      "output": {"class": "copy", "from": "sub/output"}
    }
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    losses_dict, total_loss, total_constraints = network.get_losses_initialized(with_total=True)
    print("losses:")
    pprint(losses_dict)
    assert len(losses_dict) == 1 and set(losses_dict.keys()) == {"sub/loss"}


def test_constant_layer():
  with make_scope() as session:
    config = Config()
    config.update({
      "num_outputs": 3,
      "num_inputs": 2,
      "network": {
        "output": {"class": "constant", "value": 42}
      }
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    out = network.get_default_output_layer(must_exist=True)
    v = session.run(out.output.placeholder)
    assert_equal(v.shape, ())  # (batch,), where batch==1 for broadcasting
    assert_equal(v, 42)


def test_compare_layer():
  with make_scope() as session:
    config = Config()
    config.update({
      "model": "/tmp/test-compare-layer-model",
      "num_outputs": 3,
      "num_inputs": 2,
      "network": {
        "const": {"class": "constant", "value": 3},
        "output": {"class": "compare", "from": "const", "value": 3}
      }
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    out = network.get_default_output_layer(must_exist=True)
    v = session.run(out.output.placeholder)
    assert_equal(v.shape, ())  # (batch,), where batch==1 for broadcasting
    assert_equal(v.dtype, numpy.dtype("bool"))
    assert_equal(v, True)


def test_shift_layer():
  with make_scope() as session:
    import numpy as np
    batch_size = 8
    time_size = 20
    feat_size = 10
    shift_amount = 5  # right-shift of 5 elements
    config = Config()
    config.update({
      "num_outputs": feat_size,
      "num_inputs": feat_size,
      "network": {
        "output": {"class": "shift_axis", "from": ["data"], "amount": shift_amount, "pad": True, "axis": "T"}
      }
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    out = network.get_default_output_layer(must_exist=True)
    input_data = np.ones(shape=(batch_size, time_size, feat_size))
    input_data[0, :, 0] = np.arange(time_size) # 0..time_size in time-axis
    feed_dict = {network.layers['data'].output.placeholder: input_data}
    v = session.run(out.output.placeholder, feed_dict)

    assert_equal(v.shape, (batch_size, time_size, feat_size))
    assert_equal(np.equal(v[0, shift_amount:, 0], np.arange(time_size-shift_amount)).all(), True)
    assert_equal((v[:,:shift_amount,:] == 0).all(), True)  # padding
    assert_equal((v[1:,shift_amount:,:] == 1).all(), True)


def test_ReinterpretDataLayer_change_batch_to_spatial():
  new_spatial_dim = SpatialDim("new-spatial")
  net_dict = {
    "output": {"class": "reinterpret_data", "from": "data", "set_dim_tags": {"B": new_spatial_dim}}}
  config = Config({"extern_data": {"data": {"dim": 7}}})
  with make_scope():
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)


def test_layer_base_get_out_data_from_opts():
  with make_scope() as session:
    config = Config()
    config.update({
      "num_inputs": 4,
      "num_outputs": 3
    })
    network = TFNetwork(config=config)
    input_data = network.extern_data.data["data"]
    target_data = network.extern_data.data["classes"]
    assert input_data.dim == 4
    assert input_data.shape == (None, 4)
    assert not input_data.sparse
    assert input_data.dtype == "float32"
    assert target_data.dim == 3
    assert target_data.shape == (None,)
    assert target_data.sparse
    assert target_data.dtype == "int32"
    out = LayerBase._base_get_out_data_from_opts(network=network, name="output", target="classes")
    # Output data type is a non-sparse version of the targets by default.
    assert out.dim == target_data.dim
    assert out.shape == target_data.shape_dense
    assert not out.sparse
    assert out.dtype == "float32"


def test_ReduceLayer_reduce4d():
  config = Config()
  config.update({
    "num_inputs": 4,
    "num_outputs": 3,
    "debug_print_layer_output_template": True
  })
  network = TFNetwork(config=config)
  src_layer = InternalLayer(
    name="src", network=network, output=Data(name="src", shape=(None, 4, 512), auto_create_placeholders=True))
  print("src:", src_layer)
  opts = {
    'axes': "dim:4",
    'keep_dims': True,
    'mode': 'mean',
    'name': 'c_out_reduce',
    'network': network,
    'sources': [src_layer]}
  out = ReduceLayer.get_out_data_from_opts(**opts)
  layer = ReduceLayer(output=out, **opts)
  print("layer:", layer)


def test_ReduceLayer_mean():
  net_dict = {
    "output": {"class": "reduce", "mode": "mean", "from": "data", "axis": "T"}
  }
  config = Config(dict(
    extern_data={"data": {"shape": (None, 4)}}
  ))
  with make_scope() as session:
    network = TFNetwork(config=config)
    network.construct_from_dict(net_dict)
    in_ = network.extern_data.get_default_input_data()
    out = network.get_default_output_layer().output
    assert out.batch_dim_axis == 0 and out.shape == (4,)
    in_v, seq_len, out_v = session.run(
      (in_.placeholder, in_.get_sequence_lengths(), out.placeholder),
      feed_dict=make_feed_dict(network.extern_data))
    n_batch = in_v.shape[0]
    assert n_batch == seq_len.shape[0] == out_v.shape[0]
    for b in range(n_batch):
      numpy.testing.assert_almost_equal(out_v[b], in_v[b, :seq_len[b]].mean(axis=0))


def test_reduce_repeat_1102():
  # https://github.com/rwth-i6/returnn/issues/1102
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim

  time_dim = SpatialDim('time')
  F_dim = FeatureDim('F', 1)
  speech_dim = SpatialDim('speech')
  speech_feat_dim = FeatureDim('speech-feat', 3)

  config = Config(dict(extern_data={
    'emb': {
      'dim_tags': (batch_dim, time_dim, F_dim),
      'dtype': 'float32',
      'available_for_inference': True
    },
    'durations': {
      'dim_tags': (batch_dim, time_dim),
      'dtype': 'int32',
      'available_for_inference': True
    },
    'target_speech': {
      'dim_tags': (batch_dim, speech_dim, speech_feat_dim),
      'dtype': 'float32',
      'available_for_inference': True
    }
  }))

  net_dict = {
    'nartts_model_reduce': {
      'class': 'copy',
      'from': 'reduce',
      'loss': 'as_is',
      'out_shape': {batch_dim, speech_feat_dim}
    },
    'output': {
      'class': 'copy',
      'from': 'repeat',
      'out_shape': {batch_dim, F_dim, speech_dim}
    },
    'reduce': {
      'class': 'reduce',
      'from': 'data:target_speech',
      'mode': 'mean',
      'axis': speech_dim,
      'out_shape': {batch_dim, speech_feat_dim}
    },
    'repeat': {
      'class': 'repeat',
      'from': 'data:emb',
      'repetitions': 'data:durations',
      'axis': time_dim,
      'out_dim': speech_dim,
      'out_shape': {batch_dim, F_dim, speech_dim}
    }
  }

  with make_scope() as session:
    net = TFNetwork(config=config, eval_flag=True)
    net.construct_from_dict(net_dict)

    d = net.extern_data.data
    feed_dict = {
      d["emb"].placeholder: [[[1.], [2.], [0.]]],
      d["emb"].size_placeholder[0]: [3],
      d["durations"].placeholder: [[1, 2, 1]],
      d["target_speech"].placeholder: [[[1., 2., 3.], [4., 5., 6.], [1., 2., 3.], [1., 2., 3.]]],
      d["target_speech"].size_placeholder[0]: [4],
    }
    # TODO ERROR this is not extern_data/placeholders/target_speech/target_speech_dim0_size
    #   but Tensor("repeat/Sum:0", shape=(?,), dtype=int32)
    print(d["target_speech"].size_placeholder[0])
    fetches = net.get_fetches_dict()
    session.run(fetches, feed_dict=feed_dict)


def test_SoftmaxOverSpatialLayer_start():
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData({"data": {"shape": (None, 5)}}), config=Config())
    rnd = numpy.random.RandomState(42)
    n_batch = 3
    n_time = 4
    n_dim = 7
    start_idxs = numpy.array([[3], [0], [1]]).astype("int32")  # (B, 1)
    input_np = rnd.normal(size=(n_batch, n_time, n_dim)).astype("float32")  # (B, T, D)
    src = InternalLayer(
      name="src", network=net, output=Data(**{"name": "src", "shape": (None, n_dim), "time_dim_axis": 1}))
    start = InternalLayer(
      name="start", network=net, output=Data(**{"name": "start", "shape": (1,), "dtype": "int32"}))
    start.output.placeholder = tf.constant(start_idxs)
    start.output.size_placeholder = {}
    print("input:", src.output)
    src.output.placeholder = tf.constant(input_np, dtype=tf.float32)
    src.output.size_placeholder = {0: tf.constant([n_time] * n_batch)}  # not sure if enough...
    opts = {"network": net, "name": "softmax_over_spatial_test", "sources": [src],
            "start": start, "use_time_mask": True}
    out_data = SoftmaxOverSpatialLayer.get_out_data_from_opts(**opts)
    print("output:", out_data)
    out_data.sanity_check(ignore_placeholder=True)  # placeholder might be overwritten later
    assert_equal(out_data.shape, (n_dim, None))  # layer moves time-dim to back
    layer = SoftmaxOverSpatialLayer(output=out_data, **opts)
    assert_equal(layer.output.shape, (n_dim, None))
    try:
      out_np = session.run(
        layer.output.placeholder, feed_dict={net.extern_data.get_batch_info().dim: n_batch})
    except Exception as exc:
      from returnn.tf.network import help_on_tf_exception
      help_on_tf_exception(session=session, exception=exc, fetches=layer.output.placeholder)
      raise
    assert_equal(out_np.shape, (n_batch, n_dim, n_time))
    # check if masking worked
    range_idxs = numpy.ones_like(start_idxs) * numpy.expand_dims(numpy.arange(n_time), axis=0)
    cond = range_idxs < numpy.broadcast_to(start_idxs, [n_batch, n_time])  # (B, T)
    cond = numpy.expand_dims(cond, axis=1)
    cond = numpy.broadcast_to(cond, [n_batch, n_dim, n_time])  # (B, D, T)
    assert_equal(cond.sum(), n_dim*start_idxs.sum())  # check num of conds
    numpy.testing.assert_array_equal(out_np[cond], 0)


def test_SoftmaxOverSpatialLayer_window():
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData({"data": {"shape": (None, 5)}}), config=Config())
    rnd = numpy.random.RandomState(42)
    n_batch = 4
    n_time = 9
    n_dim = 1
    window_size = 5
    window_start_idxs = numpy.array([3, 0, 1, 7]).astype("int32")  # (B,)
    seqlens = numpy.array([5, 7, 3, 9]).astype("int32")
    input_np = rnd.normal(size=(n_batch, n_time, n_dim)).astype("float32")  # (B, T, D)
    src = InternalLayer(
      name="src", network=net, output=Data(**{"name": "src", "shape": (None, n_dim), "time_dim_axis": 1}))
    window_start = InternalLayer(
      name="window_start", network=net, output=Data(**{"name": "window_start", "shape": (), "dtype": "int32"}))
    window_start.output.placeholder = tf.constant(window_start_idxs)  # (B,)
    window_start.output.size_placeholder = {}
    print("input:", src.output)
    src.output.placeholder = tf.constant(input_np, dtype=tf.float32)
    src.output.size_placeholder = {0: tf.constant(seqlens, dtype=tf.int32)}
    opts = {"network": net, "name": "softmax_over_spatial_test", "sources": [src],
            "window_start": window_start, "window_size": window_size}
    out_data = SoftmaxOverSpatialLayer.get_out_data_from_opts(**opts)
    print("output:", out_data)
    out_data.sanity_check(ignore_placeholder=True)  # placeholder might be overwritten later
    assert_equal(out_data.shape, (n_dim, None))  # layer moves time-dim to back
    layer = SoftmaxOverSpatialLayer(output=out_data, **opts)
    layer.output.sanity_check()
    assert_equal(layer.output.shape, (n_dim, None))
    out_np = session.run(layer.output.placeholder, feed_dict=make_feed_dict(net.extern_data, n_batch=n_batch))
    assert_equal(out_np.shape, (n_batch, n_dim, n_time))
    # check if window masking worked:
    # handle edge cases correctly: (start is 0-based)
    # 1. if the energy time-dim is less than `window_size`, we adjust the window size.
    # 2. for each seq, we adjust the window so that no elements after the seq-len are indexed.
    # seq[0]: start=3, seqlen=5 -> [1, 1, 1, 1, 1, 0, 0, 0, 0]
    # seq[1]: start=0, seqlen=7 -> [1, 1, 1, 1, 1, 0, 0, 0, 0]
    # seq[2]: start=1, seqlen=3 -> [1, 1, 1, 0, 0, 0, 0, 0, 0]
    # seq[3]: start=7, seqlen=9 -> [0, 0, 0, 0, 1, 1, 1, 1, 1]
    mask = numpy.array([
      [0, 0, 0, 1, 1, 0, 0, 0, 0],
      [1, 1, 1, 1, 1, 0, 0, 0, 0],
      [0, 1, 1, 0, 0, 0, 0, 0, 0],
      [0, 0, 0, 0, 1, 1, 1, 1, 1]
    ], dtype=numpy.bool)  # (B, T)
    print("mask", mask)
    mask = numpy.expand_dims(mask, axis=1)
    mask = numpy.broadcast_to(mask, [n_batch, n_dim, n_time])  # (B, D, T)
    # check if layer output sums to one for each seq:
    out_sum = numpy.sum(out_np, axis=(1, 2))
    numpy.testing.assert_allclose(out_sum, [1]*n_batch, rtol=1e-5)
    numpy.testing.assert_allclose(out_np[~mask], 0, rtol=1e-5)  # check if masking worked


def test_SplitLayer_after_SplitDimsLayer():
  n_batch, n_time, n_in = 7, 3, 40
  config = Config({
    "extern_data": {"data": {"dim": n_in}},
    "debug_print_layer_output_template": True,
  })
  with make_scope():
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "split_heads": {"class": "split_dims", "dims": (2, -1), "axis": "F", "from": "data:data"},  # [B,T,2,F|20]
      "split_qkv": {"class": "split", "size_splits": (5, 5, 10), "axis": "F", "from": "split_heads"},
      "output": {"class": "copy", "from": "split_qkv/0"}})  # [B,T,2,F|5]
    out_t = net.get_default_output_layer().output.placeholder
    assert out_t.shape.as_list() == [None, None, 2, 5]


def test_SplitLayer_search():
  n_batch, n_time, n_in, n_out = 7, 3, 10, 10
  beam_size = 4
  config = Config({
    "extern_data": {
      "data": {"dim": n_in},
      "classes": {"dim": n_out, "sparse": True, "available_for_inference": False}},
    "debug_print_layer_output_template": True
  })
  with make_scope():
    net = TFNetwork(config=config, search_flag=True, train_flag=False, eval_flag=True)
    net.construct_from_dict({
      "encoder_seq": {"class": "linear", "activation": "tanh", "n_out": 5, "from": "data:data"},
      "encoder": {"class": "reduce", "mode": "sum", "from": ["encoder_seq"], "axis": "T"},
      "output": {"class": "rec", "from": [], "target": "classes", "max_seq_len": 20, "unit": {
        "embed": {"class": "linear", "from": ["prev:output"], "activation": None, "n_out": 10},
        "split": {"class": "split", "size_splits": (5, 5), "axis": "F", "from": ["embed"]},
        "output_prob": {"class": "softmax", "from": ["split/0", "base:encoder"], "target": "classes", "loss": "ce"},
        "output": {
          "class": "choice", "target": "classes", "beam_size": beam_size, "from": ["output_prob"], "initial_output": 0},
        "end": {"class": "compare", "from": ["output"], "value": 0}
      }},
      "decision": {"class": "decide", "from": ["output"], "loss": "edit_distance", "target": "classes"}})


def test_SplitDimsLayer_simple_feat():
  n_batch, n_time, n_in = 7, 3, 20
  config = Config({
    "extern_data": {"data": {"dim": n_in}},
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {"class": "split_dims", "axis": "f", "dims": (-1, 5), "from": "data:data"}})
    out_t = net.get_default_output_layer().output.placeholder
    assert out_t.shape.as_list() == [None, None, 4, 5]
    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_time, n_in))
    out_v = session.run(out_t, feed_dict={net.extern_data.data["data"].placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, n_time, 4, 5)
    numpy.testing.assert_almost_equal(out_v, in_v.reshape(out_v.shape))


def test_SplitDimsLayer_simple_time():
  n_batch, n_time, n_in = 7, 3, 20
  config = Config({
    "extern_data": {"data": {"dim": n_in}},
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {"class": "split_dims", "axis": "t", "dims": (-1, 1), "from": "data:data"}})
    assert_equal(
      net.get_default_output_layer().output.get_dim_tag(1),
      net.extern_data.get_default_input_data().get_dim_tag(1))
    out_t = net.get_default_output_layer().output.placeholder
    assert out_t.shape.as_list() == [None, None, 1, 20]
    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_time, n_in))
    out_v = session.run(out_t, feed_dict={net.extern_data.data["data"].placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, n_time, 1, n_in)
    numpy.testing.assert_almost_equal(out_v, in_v.reshape(out_v.shape))


def test_SplitDimsLayer_simple_time2():
  n_batch, n_time, n_in = 7, 3, 20
  config = Config({
    "extern_data": {"data": {"dim": n_in}},
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {"class": "split_dims", "axis": "t", "dims": (1, -1), "from": "data:data"}})
    in_ = net.extern_data.get_default_input_data()
    out = net.get_default_output_layer().output
    print(in_)
    print(out)
    assert_equal(out.get_dim_tag(2), in_.get_dim_tag(1))
    assert out.time_dim_axis == 2
    out_t = out.placeholder
    assert out_t.shape.as_list() == [None, 1, None, 20]
    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_time, n_in))
    out_v = session.run(out_t, feed_dict={in_.placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, 1, n_time, n_in)
    numpy.testing.assert_almost_equal(out_v, in_v.reshape(out_v.shape))


def test_SplitDimsLayer_resolve_dims():
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=3 * 5, new_dims=(3, -1)), (3, 5))
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=3 * 5, new_dims=(3, 5)), (3, 5))
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=3 * 5, new_dims=(5, -1)), (5, 3))
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=2 * 3 * 5, new_dims=(-1, 3, 5)), (2, 3, 5))
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=2 * 3 * 5, new_dims=(2, -1, 5)), (2, 3, 5))
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=2 * 3 * 5, new_dims=(2, 3, -1)), (2, 3, 5))
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=2 * 3 * 5, new_dims=(2, 3, -1, 1)), (2, 3, 5, 1))

  assert_equal(SplitDimsLayer._resolve_dims(old_dim=3 * 5, new_dims=(3, -1), pad_to_multiples=True), (3, 5))
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=3 * 5 + 1, new_dims=(3, -1), pad_to_multiples=True), (3, 6))
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=2 * 3 * 5, new_dims=(2, 3, -1), pad_to_multiples=True), (2, 3, 5))
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=2 * 3 * 5, new_dims=(2, 3, -1, 1), pad_to_multiples=True), (2, 3, 5, 1))
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=2 * 3 * 5 + 2, new_dims=(2, 3, -1), pad_to_multiples=True), (2, 3, 6))
  assert_equal(SplitDimsLayer._resolve_dims(old_dim=2 * 3 * 5 + 2, new_dims=(2, 3, -1, 1), pad_to_multiples=True), (2, 3, 6, 1))


def test_SplitDimsLayer_batch_feature_major_keep_feature():
  n_batch, n_time, n_in = 7, 3, 5
  config = Config({
    "extern_data": {
      "data": {"dim": n_in, "shape": (n_in, None)}  # [B,D,T], i.e. batch-feature-major
    }
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    assert net.extern_data.get_default_input_data().is_batch_feature_major
    net.construct_from_dict({
      "output": {'class': 'split_dims', 'from': 'data', 'axis': 'T', 'dims': [-1, 1]}  # [B,D,T,1]
    })
    out = net.get_default_output_layer().output
    assert_equal(out.get_dim_tag(2), net.extern_data.get_default_input_data().get_time_dim_tag())
    assert out.dim_tags[1].dimension == n_in and out.dim_tags[3].dimension == 1
    assert out.placeholder.shape.as_list() == [None, n_in, None, 1]
    assert out.feature_dim_axis == 1  # https://github.com/rwth-i6/returnn/issues/596
    assert out.time_dim_axis == 2
    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_in, n_time))
    out_v = session.run(out.placeholder, feed_dict={net.extern_data.data["data"].placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, n_in, n_time, 1)
    numpy.testing.assert_almost_equal(out_v, in_v.reshape(out_v.shape))


def test_SplitDimsLayer_split_feature():
  n_batch, n_time, n_in = 7, 3, 5
  config = Config({
    "extern_data": {
      "data": {"dim": n_in}  # [B,T,D]
    }
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {'class': 'split_dims', 'from': 'data', 'axis': 'F', 'dims': [-1, 1]}
    })
    out = net.get_default_output_layer().output
    assert out.batch_shape == (None, None, n_in, 1)
    assert out.feature_dim_axis == 3  # https://github.com/rwth-i6/returnn/issues/704
    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_time, n_in))
    out_v = session.run(out.placeholder, feed_dict={net.extern_data.data["data"].placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, n_time, n_in, 1)
    numpy.testing.assert_almost_equal(out_v, in_v.reshape(out_v.shape))


def test_SplitDimsLayer_dim_tags():
  from returnn.tf.util.data import batch_dim
  for window_static_dim in [2, 1]:
    time_dim = SpatialDim("in-time")
    rem_dim = SpatialDim("rem-time")
    window_dim = FeatureDim("window", window_static_dim)
    feat_dim = FeatureDim("feat", 3)
    config = Config({
      "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}}})
    with make_scope():
      net = TFNetwork(config=config)
      net.construct_from_dict({
        "output": {
          'class': 'split_dims', 'from': 'data', 'axis': time_dim, 'dims': [rem_dim, window_dim],
          'out_shape': {batch_dim, rem_dim, window_dim, feat_dim}}
      })


def test_SplitDimsLayer_dim_tags_expand():
  from returnn.tf.util.data import batch_dim
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("feat", 3)
  expand_dim = SpatialDim("expand_dim", 1)
  config = Config({
    "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}}})
  with make_scope():
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {
        'class': 'split_dims', 'from': 'data', 'axis': feat_dim, 'dims': [feat_dim, expand_dim],
        'out_shape': {batch_dim, time_dim, feat_dim, expand_dim}}
    })


def test_SplitDimsLayer_dim_tags_split_batch_simple():
  # https://github.com/rwth-i6/returnn/issues/908
  # https://github.com/rwth-i6/pytorch-to-returnn/pull/78
  from returnn.tf.util.data import batch_dim
  time_dim = SpatialDim("in-time")
  feat_dim = FeatureDim("feat", 3)
  config = Config({
    "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}}})
  with make_scope():
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {
        'class': 'split_dims', 'from': 'data', 'axis': batch_dim, 'dims': [1, -1]}
    })


def test_SplitDimsLayer_dim_tags_split_batch():
  from returnn.tf.util.data import batch_dim
  time_dim = SpatialDim("in-time")
  feat_dim = FeatureDim("feat", 3)
  config = Config({
    "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}}})
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "merge": {"class": "merge_dims", "from": "data", "axes": [batch_dim, feat_dim]},
      "output": {'class': 'split_dims', 'from': "merge", 'axis': batch_dim, 'dims': [batch_dim, feat_dim]}
    })
    in_ = net.extern_data.get_default_input_data()
    assert in_.batch.is_global_batch()
    merged = net.layers["merge"].output
    out = net.layers["output"].output
    print(merged)
    assert not merged.batch.is_global_batch()
    assert len(merged.batch.virtual_dims) == 2
    b1, b2 = merged.batch.virtual_dims
    assert isinstance(b2, BatchInfo.FixedDim)
    assert b2.dim_tag == feat_dim
    print(out)
    assert in_.batch == out.batch and out.batch.is_global_batch()
    session.run(out.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_SplitDimsLayer_dyn_dim_tags_with_batch():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("feat", 3)
  config = Config({
    "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}}})
  time2_dim = time_dim * 2
  print("time:", time_dim)
  print("time * 2:", time2_dim)
  net_dict = {
    'Flatten': {'class': 'merge_dims', 'from': 'data', 'axes': ['B', 'T'], 'keep_order': True},
    'Range_Length': {'class': 'length', 'axis': 'B', 'from': 'data'},
    'Range_Length_1': {'class': 'length', 'axis': 'T', 'from': 'data'},
    'Range_Reduce': {'class': 'reduce', 'mode': 'max', 'axes': ['B'], 'from': 'Range_Length_1'},
    'Range_mul': {'class': 'combine', 'kind': 'mul', 'out_shape': set(), 'from': ['Range_Length', 'Range_Reduce']},
    'Range': {'class': 'range_from_length', 'from': 'Range_mul'},  # [T]
    'Cat': {'class': 'concat', 'from': [('Range', 'T'), ('Range', 'T')]},
    'Unflatten__Unflatten_mul_unnamed_const': {'class': 'constant', 'value': 2},
    'Unflatten__Unflatten_mul': {
      'class': 'combine', 'kind': 'mul', 'out_shape': set(),
      'from': ['Range_Reduce', 'Unflatten__Unflatten_mul_unnamed_const']},
    'output': {
      'class': 'split_dims', 'from': 'Cat', 'axis': 'T', 'dims': [batch_dim, time2_dim]}
  }
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    in_ = net.extern_data.get_default_input_data()
    out = net.get_default_output_layer().output
    assert out.dim_tags == (batch_dim, time2_dim)
    x, x_size, y, y_size = session.run(
      (in_.placeholder, in_.size_placeholder.as_dict(), out.placeholder, out.size_placeholder.as_dict()),
      feed_dict=make_feed_dict(net.extern_data))
    assert isinstance(x, numpy.ndarray)
    assert isinstance(x_size, dict) and set(x_size.keys()) == {0}
    x_size = x_size[0]
    assert isinstance(x_size, numpy.ndarray)
    n_batch, n_time, n_feat = x.shape
    assert x_size.shape == (n_batch,) and max(x_size) == n_time
    assert isinstance(y, numpy.ndarray)
    assert isinstance(y_size, dict) and set(y_size.keys()) == {0}
    y_size = y_size[0]
    assert isinstance(y_size, numpy.ndarray)
    assert y.shape == (n_batch, n_time * 2)
    assert y_size.shape == (n_batch,) and max(y_size) == n_time * 2


def test_ReshapeLayer():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  n_batch, n_time, n_dim = 2, 4, 3
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("feature", dimension=n_dim)
  config = Config({
    "extern_data": {
      "data": {"dim_tags": [batch_dim, time_dim, feat_dim]}  # [B,T,D]
    }
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {
        'class': 'reshape', 'from': 'data',
        'in_dims': [time_dim, feat_dim],
        'out_dims': [feat_dim, time_dim]
      }
    })
    out = net.get_default_output_layer().output
    assert out.dim_tags == (batch_dim, feat_dim, time_dim)
    in_v = numpy.arange(n_time, dtype="float32")[None, :, None] + numpy.zeros((n_batch, n_time, n_dim), dtype="float32")
    feed_dict = {
      net.extern_data.data["data"].placeholder: in_v,
      net.extern_data.data["data"].size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
      net.extern_data.get_batch_info().dim: n_batch,
    }
    out_v = session.run(out.placeholder, feed_dict=feed_dict)
    print(out_v.shape)
    print(out_v)
    assert out_v.shape == (n_batch, n_dim, n_time)
    numpy.testing.assert_array_equal(in_v.reshape((n_batch, n_dim, n_time)), out_v)


def test_out_shape():
  # https://github.com/rwth-i6/returnn/issues/706
  # Note: Using SplitDimsLayer would also be nice to test out_shape. Or any layer which creates a new dim.
  # However, for that, we need https://github.com/rwth-i6/returnn/issues/597 first.
  from returnn.tf.util.data import batch_dim, VerifyOutShapeException
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("feature", dimension=10)
  config = Config({
    "extern_data": {
      "data": {"dim_tags": [batch_dim, time_dim, feat_dim]}  # [B,T,D]
    }
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {
        'class': 'softmax_over_spatial', 'from': 'data',
        "out_shape": {batch_dim, time_dim, feat_dim}
      }
    })
    out = net.get_default_output_layer().output
    session.run(out.placeholder, feed_dict=make_feed_dict(net.extern_data))
  with make_scope():
    other_feat_dim = FeatureDim("other-feature", dimension=10)
    net = TFNetwork(config=config)
    # noinspection PyBroadException
    try:
      net.construct_from_dict({
        "output": {
          'class': 'softmax_over_spatial', 'from': 'data',
          "out_shape": {batch_dim, time_dim, other_feat_dim}
        }
      })
    except VerifyOutShapeException as exc:
      print("Got expected exception: %r" % exc)
    else:
      raise Exception("Expected an exception but did not get any.")


def _check_MergeDimsLayer(session, in_data_opts, in_static_shape, opts, out_data_shape, out_static_shape,
                          in_sizes=None, out_sizes=None):
  """
  :param tf.compat.v1.Session session:
  :param dict[str] in_data_opts:
  :param tuple[int] in_static_shape:
  :param dict[str] opts: for MergeDimsLayer
  :param tuple[int|None] out_data_shape:
  :param tuple[int] out_static_shape:
  :param dict[int,tuple[int]]|None in_sizes:
  :param dict[int,tuple[int]]|None out_sizes:
  :rtype: MergeDimsLayer
  """
  net = TFNetwork(extern_data=ExternData())
  rnd = numpy.random.RandomState(42)
  src = InternalLayer(name="src", network=net, output=Data(name="src", **in_data_opts))
  print("input:", src.output)
  src.output.placeholder = tf.constant(rnd.normal(size=in_static_shape).astype("float32"), dtype=tf.float32)
  src.output.size_placeholder = {}
  if src.output.batch_dim_axis is not None:
    n_batch = in_static_shape[src.output.batch_dim_axis]
    for axis, dim in enumerate(src.output.batch_shape):
      axis_wo_b = src.output.get_batch_axis_excluding_batch(axis)
      if dim is None and axis_wo_b is not None:
        if in_sizes and axis_wo_b in in_sizes:
          src.output.size_placeholder[axis_wo_b] = tf.constant(in_sizes[axis_wo_b])
        else:
          src.output.size_placeholder[axis_wo_b] = tf.fill([n_batch], in_static_shape[axis])

  opts = opts.copy()
  print("opts:", opts)
  opts.update({"network": net, "name": "merge_dims_test", "sources": [src]})
  out_data = MergeDimsLayer.get_out_data_from_opts(**opts)
  out_data.sanity_check(ignore_placeholder=True)  # placeholder might be overwritten later
  assert_equal(out_data.shape, out_data_shape)
  layer = MergeDimsLayer(output=out_data, **opts)
  assert_equal(layer.output.shape, out_data_shape)
  out_np, size_placeholder = session.run([layer.output.placeholder, layer.output.size_placeholder.as_dict()])
  print("output:", out_data)
  assert_equal(out_np.shape, out_static_shape)

  if out_sizes:
    assert_equal(sorted(size_placeholder.keys()), sorted(out_sizes))
    for k in size_placeholder.keys():
      numpy.testing.assert_array_equal(size_placeholder[k], out_sizes[k])

  return layer


def test_MergeDimsLayer_basic():
  with make_scope() as session:
    _check_MergeDimsLayer(session, {"shape": (4, 7), "time_dim_axis": None}, (2, 4, 7), {"axes": ["dim:4", "dim:7"]}, (4 * 7,), (2, 4 * 7))
    _check_MergeDimsLayer(session, {"shape": (4, None, 7), "time_dim_axis": None}, (2, 4, 3, 7), {"axes": ["dim:4", "dim:7"]}, (None, 4 * 7), (2, 3, 4 * 7))
    _check_MergeDimsLayer(session, {"shape": (4, None, 7), "time_dim_axis": 2}, (2, 4, 3, 7), {"axes": ["dim:4", "dim:7"]}, (None, 4 * 7), (2, 3, 4 * 7))
    _check_MergeDimsLayer(session, {"shape": (1, None), "time_dim_axis": 2, "feature_dim_axis": 1}, (2, 1, 4), {"axes": ["F", "T"]}, (None,), (2, 4))


def test_MergeDimsLayer_size_placeholder():
  with make_scope() as session:
    _check_MergeDimsLayer(
      session,
      {"shape": (None, 2), "time_dim_axis": 1, "feature_dim_axis": 2}, (3, 4, 2), {"axes": ["T", "F"]}, (None,), (3, 8),
      in_sizes={0: (4, 2, 1)}, out_sizes={0: (8, 4, 2)})


def test_MergeDimsLayer_batch_time_ext():
  with make_scope() as session:
    n_batch = 11
    n_time = 13
    _check_MergeDimsLayer(
      session, {"shape": (None, 5, 3)}, (n_batch, n_time, 5, 3), {"axes": ["B", "T"]}, (5, 3), (n_batch * n_time, 5, 3))


def test_MergeDimsLayer_batch_time_time_major():
  with make_scope() as session:
    n_batch = 11
    n_time = 13
    layer = _check_MergeDimsLayer(
      session,
      {"shape": (None, 5), "time_dim_axis": 0, "batch_dim_axis": 1}, (n_time, n_batch, 5),
      {"axes": ["B", "T"]}, (5,), (n_time * n_batch, 5))
    assert layer.output.batch_dim_axis == 0
    assert layer.output.time_dim_axis is None


def test_MergeDimsLayer_batch_time_time_major_ext():
  with make_scope() as session:
    n_batch = 11
    n_time = 13
    layer = _check_MergeDimsLayer(
      session,
      {"shape": (None, 5, 3), "time_dim_axis": 0, "batch_dim_axis": 1}, (n_time, n_batch, 5, 3),
      {"axes": ["B", "T"]}, (5, 3), (n_time * n_batch, 5, 3))
    assert layer.output.batch_dim_axis == 0
    assert layer.output.time_dim_axis is None  # Note: This behavior was changed.


def test_MergeDimsLayer_except_time_ext():
  with make_scope() as session:
    n_batch = 11
    n_time = 13
    layer = _check_MergeDimsLayer(
      session,
      {"shape": (3, None, 5), "time_dim_axis": 2}, (n_batch, 3, n_time, 5),
      {"axes": ["dim:3", "dim:5"]}, (None, 15), (n_batch, n_time, 15))
    assert layer.output.batch_dim_axis == 0 and layer.output.time_dim_axis == 1


def test_MergeDimsLayer_static_time():
  with make_scope() as session:
    n_batch = 11
    layer = _check_MergeDimsLayer(
      session,
      {"shape": (3, 5), "time_dim_axis": 1}, (n_batch, 3, 5),
      {"axes": ["dim:3", "dim:5"]}, (15,), (n_batch, 15))
    assert layer.output.batch_dim_axis == 0 and layer.output.feature_dim_axis == 1
    assert layer.output.time_dim_axis is None


def test_MergeDimsLayer_feat_static_static():
  with make_scope() as session:
    n_batch = 11
    layer = _check_MergeDimsLayer(
      session,
      {"shape": (None, 8, 2, 3), "feature_dim_axis": 2}, (n_batch, 7, 8, 2, 3),
      {"axes": ["F", "dim:2"]}, (None, 16, 3), (n_batch, 7, 16, 3))
    assert (layer.output.batch_dim_axis, layer.output.time_dim_axis, layer.output.feature_dim_axis) == (0, 1, 2)


def test_MergeDimsLayer_dim_tags():
  n_batch = 3
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData())
    rnd = numpy.random.RandomState(42)

    src_data = Data("input", shape=(None, None, None, 1), feature_dim_axis=None)
    input_static_shape = (n_batch, 7, 1, 2, 1)
    src_data.placeholder = tf.constant(rnd.normal(size=input_static_shape).astype("float32"), dtype=tf.float32)
    src_data.size_placeholder = {}
    from returnn.tf.util.basic import Dim
    # map axis_wo_batch -> (tag description, dyn_size)
    tag_names_with_dyn_size = {0: ("key-chunk", [4, 2, 3]), 1: ("key-window", [1, 1, 1]), 2: ("att-heads", [2, 2, 2])}
    for axis_wo_batch, (description, dyn_size) in tag_names_with_dyn_size.items():
      tag = Dim(description=description, kind=Dim.Types.Spatial)
      dyn_size = tf.constant(dyn_size)
      tag.set_tag_on_size_tensor(dyn_size)
      src_data.size_placeholder[axis_wo_batch] = dyn_size
    print('in data:', src_data)  # should be [B,T|'key-chunk',1|'key-window',2|'att-heads',1]
    assert (
      src_data.get_axis_by_tag_name('key-chunk') == 1 and src_data.get_axis_by_tag_name('key-window') == 2 and
      src_data.get_axis_by_tag_name('att-heads') == 3)

    merge_axes = ['stag:key-window', 'dim:1']
    print('merge axes:', merge_axes)

    src = InternalLayer(name="src", network=net, output=src_data)
    opts = {"network": net, "name": "merge_dims_test", "sources": [src], "axes": merge_axes}
    out_data = MergeDimsLayer.get_out_data_from_opts(**opts)
    out_data.sanity_check(ignore_placeholder=True)  # placeholder might be overwritten later
    print('template out data:', out_data)  # should be [B,T|'key-chunk',1|<anything>,2|'att-heads',1]
    assert out_data.shape == src_data.shape[:-1]
    assert out_data.get_axis_by_tag_name('key-chunk') == 1 and out_data.get_axis_by_tag_name('att-heads') == 3

    layer = MergeDimsLayer(output=out_data, **opts)
    layer.output.sanity_check()
    out_data = layer.output
    print('layer out data:', out_data)
    assert out_data.shape == src_data.shape[:-1]
    assert out_data.get_axis_by_tag_name('key-chunk') == 1 and out_data.get_axis_by_tag_name('att-heads') == 3


def test_MergeDimsLayer_SplitBatchTimeLayer_time_major():
  n_batch = 3
  n_time = 4
  n_input_dim = 5
  # Time major
  input_data = numpy.arange(n_time * n_batch * n_input_dim).reshape((n_time, n_batch, n_input_dim)).astype("float32")
  with make_scope() as session:
    seq_lens = [n_time, n_time - 1, n_time - 2]
    assert len(seq_lens) == n_batch and all([s > 0 for s in seq_lens])
    net = TFNetwork(extern_data=ExternData(), config=Config({"debug_print_layer_output_template": True}))
    input_layer = net.add_layer(
      "input", InternalLayer,
      output=Data(
        name="input", shape=(None, n_input_dim), time_dim_axis=0, batch_dim_axis=1,
        placeholder=tf.constant(input_data),
        size_placeholder={0: tf.constant(seq_lens)}))
    assert input_layer.output.is_time_major
    net.construct_from_dict({
      "merge_dims": {"class": "merge_dims", "from": "input", "axes": ["B", "T"]},
      "split_dims": {"class": "split_batch_time", "from": "merge_dims", "base": "input"},
      "output": {"class": "copy", "from": "split_dims"}
    })
    output = net.get_default_output_layer().output
    # Depending on implementation, output could be batch-major or time-major.
    output = output.copy_as_time_major()  # such that we can compare easily to input_data
    assert output.is_time_major and output.shape == (None, n_input_dim)
    output_data = session.run(output.placeholder)
    numpy.testing.assert_almost_equal(input_data, output_data)


def test_MergeDimsLayer_SplitBatchTimeLayer_two_time_axes():
  n_dim = 11
  with make_scope() as session:
    net = TFNetwork(
      config=Config({
        "extern_data": {"data": {"shape": (None, None, n_dim)}},
        "debug_print_layer_output_template": True}))
    feed_dict = make_feed_dict(net.extern_data)
    net.construct_from_dict({
      "merge_dims": {"class": "merge_dims", "from": "data", "axes": ["B", "T"]},
      "split_dims": {"class": "split_batch_time", "from": "merge_dims", "base": "data"},
      "output": {"class": "copy", "from": "split_dims"}
    })
    input_data = net.extern_data.get_default_input_data()
    print("input_data:", input_data)
    assert set(input_data.size_placeholder.keys()) == {0, 1}
    assert input_data.size_placeholder[0].name != input_data.size_placeholder[1].name
    assert input_data.get_size_dim_tag(0) != input_data.get_size_dim_tag(1)
    merged_data = net.layers["merge_dims"].output
    print("merged_data:", merged_data)
    assert set(merged_data.size_placeholder.keys()) == {0}
    assert merged_data.get_size_dim_tag(0) != input_data.get_size_dim_tag(0)
    assert merged_data.get_size_dim_tag(0) == input_data.get_size_dim_tag(1)  # like beam-search, still same dim-tag
    assert merged_data.size_placeholder[0] is not input_data.size_placeholder[1]  # but different sizes
    output_data = net.get_default_output_layer().output
    output_data = output_data.copy_as_batch_major()
    print("output_data:", output_data)
    assert output_data.shape == (None, None, n_dim)
    assert output_data.get_size_dim_tag(0) == input_data.get_size_dim_tag(0)
    assert output_data.get_size_dim_tag(1) == input_data.get_size_dim_tag(1)
    input_value = session.run(input_data.placeholder, feed_dict=feed_dict)
    merged_value = session.run(merged_data.placeholder, feed_dict=feed_dict)
    output_value = session.run(output_data.placeholder, feed_dict=feed_dict)
    assert input_value.shape == output_value.shape
    assert input_value.shape[-1] == n_dim
    print("input_value.shape:", input_value.shape)
    n_batch, n_time0, n_time1, _ = input_value.shape
    numpy.testing.assert_almost_equal(input_value, output_value)
    print("merged_value.shape:", merged_value.shape)
    assert merged_value.shape == (n_batch * n_time0, n_time1, n_dim)
    numpy.testing.assert_almost_equal(input_value, merged_value.reshape(input_value.shape))
    merged_size = session.run(merged_data.size_placeholder[0], feed_dict=feed_dict)
    input_size0, input_size1 = session.run(
      (input_data.size_placeholder[0], input_data.size_placeholder[1]), feed_dict=feed_dict)
    assert input_size0.shape == input_size1.shape == (n_batch,)
    assert merged_size.shape == (n_batch * n_time0,)
    merged_size = merged_size.reshape(n_batch, n_time0)
    assert (merged_size == input_size1[:, None]).all()


def test_MergeDimsLayer_simple_feat():
  n_batch, n_time, n_in1, n_in2 = 7, 3, 10, 32
  config = Config({
    "extern_data": {"data": {"shape": (None, n_in1, n_in2)}},
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {"class": "merge_dims", "axes": ["dim:%i" % n_in1, "dim:%i" % n_in2], "from": "data:data"}})
    out_t = net.get_default_output_layer().output.placeholder
    assert out_t.shape.as_list() == [None, None, n_in1 * n_in2]
    in_v = numpy.arange(0, n_batch * n_time * n_in1 * n_in2).astype("float32").reshape((n_batch, n_time, n_in1, n_in2))
    out_v = session.run(out_t, feed_dict={net.extern_data.data["data"].placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, n_time, n_in1 * n_in2)
    numpy.testing.assert_almost_equal(out_v, in_v.reshape(out_v.shape))


def test_MergeDimsLayer_2d_dynamic_merge_axis():
  # https://github.com/rwth-i6/returnn/issues/662
  from returnn.tf.util.data import batch_dim, Dim, ImplicitDynSizeDim
  time_dim = SpatialDim("T")
  feat_dim = FeatureDim("F", dimension=5)
  with make_scope() as session:
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}}})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "start0": {"class": "range_in_axis", "from": "data", "axis": "b", "out_shape": {batch_dim}},
      "start1": {"class": "range_in_axis", "from": "data", "axis": "t", "out_shape": {time_dim, ImplicitDynSizeDim(batch_dim)}},
      "start": {"class": "combine", "from": ["start0", "start1"], "kind": "add", "out_shape": {batch_dim, time_dim}},
      "slices": {"class": "slice_nd", "from": "data", "start": "start", "size": None},  # [B,T[B],slice[B,T],D]
      "output": {"class": "merge_dims", "from": "slices", "axes": ["f", "stag:slice"]}  # [B,T[B],merge[B,T]]
    })
    slices_layer = net.get_layer("slices")
    assert isinstance(slices_layer, SliceNdLayer)
    merge_layer = net.get_layer("output")
    assert isinstance(merge_layer, MergeDimsLayer)
    assert slices_layer.output.dim_tags[:2] == merge_layer.output.dim_tags[:2]
    assert merge_layer.output.dim_tags[0].is_batch_dim()
    assert merge_layer.output.dim_tags[1] == net.extern_data.get_default_input_data().get_time_dim_tag()
    assert merge_layer.output.dim_tags[2].dyn_size_ext.dim_tags == merge_layer.output.dim_tags[:2]  # [B,T]
    assert slices_layer.output.dim_tags[2].dyn_size_ext.dim_tags == slices_layer.output.dim_tags[:2]  # [B,T]
    assert slices_layer.output.dim_tags[1].dyn_size_ext.dim_tags == slices_layer.output.dim_tags[:1]  # [B]
    out, merged_time, sliced_time = session.run(
      (merge_layer.output.placeholder,
       merge_layer.output.dim_tags[2].dyn_size_ext.placeholder,
       slices_layer.output.dim_tags[2].dyn_size_ext.placeholder),
      make_feed_dict(net.extern_data))
    assert isinstance(out, numpy.ndarray)
    assert isinstance(merged_time, numpy.ndarray) and isinstance(sliced_time, numpy.ndarray)
    assert out.shape == merged_time.shape + (numpy.max(merged_time),)
    assert sliced_time.shape == merged_time.shape
    assert numpy.all(sliced_time * feat_dim.dimension == merged_time)


def test_MergeDimsLayer_modified_time_dim():
  n_batch, n_time, n_in = 3, 7, 2
  config = Config({
    "extern_data": {"data": {"dim": n_in}},
    "debug_print_layer_output_template": True,
    "behavior_version": 12,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "conv": {
        "class": "conv", "n_out": n_in, "filter_size": (3,), "strides": (2,), "padding": "valid", "from": "data:data"},
      "output": {"class": "merge_dims", "from": "conv", "axes": ["B", "T"], "keep_order": True},
    })
    out = net.get_default_output_layer().output
    net.initialize_params(session)
    session.run(out.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_MergeDimsLayer_unspecified_out_dim():
  # https://github.com/rwth-i6/returnn/issues/955
  # https://github.com/rwth-i6/returnn_common/issues/117
  config = Config({
    "extern_data": {"data": {"shape": (None, 3, 5)}},
  })
  out_dim = SpatialDim("out")
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {
        "class": "merge_dims", "from": "data", "axes": ["dim:3", "dim:5"], "keep_order": True,
        "out_dim": out_dim},
    })


def test_FlattenBatchLayer():
  from returnn.tf.util.data import BatchInfo
  n_batch, n_time, n_in = 3, 4, 2
  config = Config({
    "extern_data": {"data": {"dim": n_in, "dtype": "int32"}},
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {"class": "flatten_batch", "batch_major": False, "from": "data:data"}})
    in_data = net.extern_data.data["data"]
    out_data = net.get_default_output_layer().output
    assert out_data.batch_shape == (None, n_in) and not out_data.size_placeholder
    assert len(out_data.batch.virtual_dims) == 2
    batch_flat_dim0, batch_flat_dim1 = out_data.batch.virtual_dims
    assert isinstance(batch_flat_dim0, BatchInfo.PackedDim)
    assert isinstance(batch_flat_dim1, BatchInfo.GlobalBatchDim)
    assert batch_flat_dim0.sizes is in_data.size_placeholder[0]
    out_t = net.get_default_output_layer().output.placeholder
    assert out_t.shape.as_list() == [None, n_in]
    in_v = numpy.arange(0, n_batch * n_time * n_in).reshape((n_time, n_batch, n_in)).transpose(1, 0, 2)
    in_seq_lens = [4, 3, 2]
    out_v = session.run(out_t, feed_dict={
      in_data.placeholder: in_v,
      in_data.size_placeholder[0]: in_seq_lens})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (sum(in_seq_lens), n_in)
    assert_equal(out_v.tolist(), [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13], [14, 15], [18, 19]])


def test_UnflattenBatchLayer():
  n_batch, n_time, n_in = 3, 4, 2
  config = Config({
    "extern_data": {"data": {"dim": n_in, "dtype": "int32"}},
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "flatten": {"class": "flatten_batch", "from": "data:data"},
      "output": {"class": "unflatten_batch", "from": "flatten"}})
    in_data = net.extern_data.data["data"]
    out_data = net.get_default_output_layer().output
    assert out_data.batch_shape == (None, None, n_in) and out_data.size_placeholder is not None
    out_t = net.get_default_output_layer().output.placeholder
    assert out_t.shape.as_list() == [None, None, n_in]
    in_v = numpy.arange(0, n_batch * n_time * n_in).reshape((n_time, n_batch, n_in)).transpose(1, 0, 2)
    in_seq_lens = [4, 3, 2]
    out_v = session.run(out_t, feed_dict={
      in_data.placeholder: in_v,
      in_data.size_placeholder[0]: in_seq_lens})
    for b, seq_len in enumerate(in_seq_lens):
      in_v[b, seq_len:, :] = 0
    assert isinstance(out_v, numpy.ndarray)
    numpy.testing.assert_equal(out_v, in_v)


def test_UnflattenBatchLayer_time_major():
  n_batch, n_time, n_in = 3, 4, 2
  config = Config({
    "extern_data": {"data": {"dim": n_in, "dtype": "int32"}},
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "flatten": {"class": "flatten_batch", "batch_major": False, "from": "data:data"},
      "output": {"class": "unflatten_batch", "from": "flatten"}})
    in_data = net.extern_data.data["data"]
    out_data = net.get_default_output_layer().output
    assert out_data.batch_shape == (None, None, n_in) and out_data.size_placeholder is not None
    out_t = net.get_default_output_layer().output.placeholder
    assert out_t.shape.as_list() == [None, None, n_in]
    in_v = numpy.arange(0, n_batch * n_time * n_in).reshape((n_time, n_batch, n_in)).transpose(1, 0, 2)
    in_seq_lens = [4, 3, 2]
    out_v = session.run(out_t, feed_dict={
      in_data.placeholder: in_v,
      in_data.size_placeholder[0]: in_seq_lens})
    for b, seq_len in enumerate(in_seq_lens):
      in_v[b, seq_len:, :] = 0
    assert isinstance(out_v, numpy.ndarray)
    numpy.testing.assert_equal(out_v, in_v.transpose(1, 0, 2))


def test_SwitchLayer_const_no_time():
  config = Config({
    "extern_data": {
      "data": {"dim": 3, "sparse": True, "shape": ()},
    },
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "input_eq_0": {"class": "compare", "from": "data", "value": 0},  # (B,)
      "const0": {"class": "constant", "value": 0},
      "const1": {"class": "constant", "value": 1},
      "switch": {"class": "switch", "condition": "input_eq_0", "true_from": "const1", "false_from": "const0"},
      "output": {"class": "copy", "from": "switch"}})
    net.print_network_info()
    feed_dict = make_feed_dict(net.extern_data.data.values())
    out = session.run(net.get_default_output_layer().output.placeholder, feed_dict=feed_dict)
    print(out)


def test_SwitchLayer_const():
  config = Config({
    "extern_data": {
      "data": {"dim": 3, "sparse": True},
    },
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "input_eq_0": {"class": "compare", "from": "data", "value": 0},  # (B,T)
      "const0": {"class": "constant", "value": 0},
      "const1": {"class": "constant", "value": 1},
      "switch": {
        "class": "switch", "condition": "input_eq_0", "true_from": "const1", "false_from": "const0"
        },
      "output": {"class": "copy", "from": "switch"}})
    net.print_network_info()
    feed_dict = make_feed_dict(net.extern_data.data.values())
    out = session.run(net.get_default_output_layer().output.placeholder, feed_dict=feed_dict)
    print(out)


def test_SwitchLayer_masking():
  config = Config({
    "extern_data": {
      "data": {"dim": 3, "sparse": False},
    },
    "debug_print_layer_output_template": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "projected": {"class": "gather", "from": "data", "axis": "F", "position": 0}, # (B,T)
      "mask": {"class": "compare", "from": "projected", "value": 0, "kind": "greater"},  # (B,T)
      "switch": {
        "class": "switch", "condition": "mask", "true_from": "data", "false_from": float("-inf")
        },
      "output": {"class": "copy", "from": "switch"}})
    net.print_network_info()
    feed_dict = make_feed_dict(net.extern_data.data.values())
    out = session.run(net.get_default_output_layer().output.placeholder, feed_dict=feed_dict)
    print(out)


def test_SwitchLayer_template_const_from():
  net = TFNetwork(extern_data=ExternData())
  batch_dim = Dim(kind=Dim.Types.Batch, description="batch")
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("feature", dimension=2)
  # [T]
  condition = InternalLayer(network=net, name="condition", output=Data(
    "condition_output", time_dim_axis=0, feature_dim_axis=None, dim_tags=[time_dim]))
  true_from = 42
  # [B,F|2,T]
  false_from = InternalLayer(network=net, name="false_from", output=Data(
    "false_from_output", batch_dim_axis=0, time_dim_axis=2, feature_dim_axis=1,
    dim_tags=[batch_dim, feat_dim, time_dim]))

  # should be [B,F|2,T]
  switch = SwitchLayer.get_out_data_from_opts(
    'switch', condition=condition, true_from=true_from, false_from=false_from)
  assert switch.batch_ndim == 3
  assert switch.batch_dim_axis == 0 and switch.time_dim_axis == 2 and switch.feature_dim_axis == 1
  assert switch.dim == 2


def test_TopKLayer_single_axis():
  config = Config({
    "extern_data": {"data": {"shape": (None, 5)}},
  })
  with make_scope() as session:
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "output": {
        "class": "top_k", "from": "data", "axis": "F", "k": 2},  # (B,T)
    })
    in_data = net.extern_data.data["data"]
    in_np = numpy.arange(2 * 3 * 5).reshape((2, 3, 5)).astype(in_data.dtype)
    in_sizes = [3, 2]
    values = net.get_layer("output").output
    indices = net.get_layer("output/indices").output
    values_np, indices_np = session.run(
      (values.placeholder, indices.placeholder),
      feed_dict={in_data.placeholder: in_np, in_data.get_sequence_lengths(): in_sizes})
    print("inputs:\n", in_np)
    print("values:\n", values_np)
    print("indices:\n", indices_np)
    assert (indices_np == numpy.array([4, 3])[None, None]).all()


def test_TopKLayer_two_axes():
  config = Config({
    "extern_data": {"data": {"shape": (3, 5)}},
  })
  with make_scope() as session:
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "output": {
        "class": "top_k", "from": "data", "axis": ("dim:3", "dim:5"), "k": 2},  # (B,T)
    })
    in_data = net.extern_data.data["data"]
    in_np = numpy.arange(2 * 3 * 5).reshape((2, 3, 5)).astype(in_data.dtype)
    values = net.get_layer("output").output
    indices0 = net.get_layer("output/indices0").output
    indices1 = net.get_layer("output/indices1").output
    values_np, indices0_np, indices1_np = session.run(
      (values.placeholder, indices0.placeholder, indices1.placeholder),
      feed_dict={in_data.placeholder: in_np})
    print("inputs:\n", in_np)
    print("values:\n", values_np)
    print("indices0:\n", indices0_np)
    print("indices1:\n", indices1_np)
    assert (indices0_np == 2).all()
    assert (indices1_np == numpy.array([4, 3])[None]).all()


def test_TopKLayer_in_cond_kdim():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim

  time_dim = SpatialDim('time')
  feat_dim = FeatureDim('feat', 5)

  config = Config(dict(
    extern_data={
      'data': {
        'dim_tags': (batch_dim, time_dim, feat_dim),
        'dtype': 'float32',
        'available_for_inference': True
      }
    },
    debug_runtime_sanity_checks=True,
    debug_print_layer_output_shape=True,
    debug_print_layer_output=True,
  ))

  k_dim = SpatialDim('feature_masking:num')

  net_dict = {
    'specaugment_v2': {
      'class': 'subnetwork',
      'from': [],
      'subnetwork': {
        'train_flag': {'class': 'train_flag'},
        'cond': {
          'class': 'cond',
          'from': [],
          'condition': 'train_flag',
          'true_layer': {
            'class': 'subnetwork',
            'from': [],
            'subnetwork': {
              'k': {
                'class': 'random', "shape": (), "distribution": "uniform", "dtype": "int32",
                "minval": 2, "maxval": 4,
              },
              'scores': {
                "class": "reduce", "from": "base:base:data", "mode": "max", "axis": time_dim,
                "out_shape": {batch_dim, feat_dim}
              },
              'top_k': {
                'class': 'top_k',
                'from': 'scores',
                'axis': feat_dim,
                'k': 'k',
                'k_dim': k_dim,
                'sorted': True,
                'out_shape': {batch_dim, k_dim}
              },
              'range_in_axis': {
                'class': 'range_in_axis',
                'from': 'top_k/indices',
                'axis': k_dim,
                'out_shape': {k_dim}
              },
              'output': {
                'class': 'reduce', 'from': 'range_in_axis', "mode": "max",
                "axis": k_dim
              }
            }
          },
          'false_layer': {
            'class': 'subnetwork',
            'from': [],
            'subnetwork': {
              'output': {
                'class': 'constant', "value": -1, "dtype": "int32",
                'out_shape': {}
              }
            }
          },
          'name_scope': ''
        },
        "output": {"class": "copy", "from": "cond"},
      },
    },
    'output': {
      'class': 'copy',
      'from': 'specaugment_v2',
    }
  }

  with make_scope() as session:
    train_flag = tf_util.get_global_train_flag_placeholder()
    net = TFNetwork(config=config, train_flag=train_flag)
    net.construct_from_dict(net_dict)
    net.initialize_params(session)
    feed_dict = make_feed_dict(net.extern_data)
    feed_dict[train_flag] = False
    res_eval = session.run(net.get_default_output_layer().output.placeholder, feed_dict=feed_dict)
    print("eval:", res_eval)
    feed_dict[train_flag] = True
    res_train = session.run(net.get_default_output_layer().output.placeholder, feed_dict=feed_dict)
    print("train:", res_train)


def test_specaugment_pure_returnn_reduced_with_cond():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim, ImplicitSparseDim

  time_dim = SpatialDim('time')
  feat_dim = FeatureDim('feat', 50)

  config = Config(dict(
    extern_data={
      'data': {
        'dim_tags': (batch_dim, time_dim, feat_dim),
        'dtype': 'float32',
        'available_for_inference': True
      }
    },
    debug_runtime_sanity_checks=True,
    debug_print_layer_output_shape=True,
  ))

  specaugment_v2_cond_true_feature_masking_top_k_k_dim = SpatialDim('feature_masking:num')

  net_dict = {
    'specaugment_v2': {
      'class': 'subnetwork',
      'from': [],
      'subnetwork': {
        'ones': {'class': 'constant', 'value': True, 'shape': (), 'dtype': 'bool'},
        'cond': {
          'class': 'cond',
          'from': [],
          'condition': 'ones',
          'true_layer': {
            'class': 'subnetwork',
            'from': [],
            'subnetwork': {
              'dim_value': {
                'class': 'subnetwork',
                'from': [],
                'subnetwork': {
                  'length': {
                    'class': 'length',
                    'from': 'base:base:base:data:data',
                    'axis': time_dim,
                    'out_shape': {batch_dim}
                  },
                  'reduce': {
                    'class': 'reduce',
                    'from': 'length',
                    'mode': 'max',
                    'axis': (batch_dim,),
                    'out_shape': {}
                  },
                  'output': {
                    'class': 'copy',
                    'from': 'reduce',
                    'out_shape': {}
                  }
                },
                'out_shape': {}
              },
              'constant': {'class': 'constant', 'value': 2},
              'minimum': {
                'class': 'combine',
                'from': ['constant', 'dim_value'],
                'kind': 'minimum',
                'out_shape': {}
              },
              'constant_0': {'class': 'constant', 'value': 100},
              'floordiv': {
                'class': 'combine',
                'from': ['dim_value', 'constant_0'],
                'kind': 'floordiv',
                'out_shape': {}
              },
              'constant_1': {'class': 'constant', 'value': 2},
              'maximum': {
                'class': 'combine',
                'from': ['floordiv', 'constant_1'],
                'kind': 'maximum',
                'out_shape': {}
              },
              'constant_2': {'class': 'constant', 'value': 4},
              'mul': {
                'class': 'combine',
                'from': ['maximum', 'constant_2'],
                'kind': 'mul',
                'out_shape': {}
              },
              'minimum_0': {
                'class': 'combine',
                'from': ['mul', 'dim_value'],
                'kind': 'minimum',
                'out_shape': {}
              },
              'feature_masking': {
                'class': 'subnetwork',
                'from': [],
                'subnetwork': {
                  'random': {
                    'class': 'subnetwork',
                    'from': [],
                    'subnetwork': {
                      'random': {
                        'class': 'random',
                        'shape': [
                          batch_dim
                        ],
                        'distribution': 'uniform',
                        'minval': 2,
                        'maxval': 6,
                        'dtype': 'int32',
                      },
                      'output': {
                        'class': 'copy',
                        'from': 'random',
                        'out_shape': {batch_dim}
                      }
                    },
                    'out_shape': {batch_dim}
                  },
                  'random_0': {
                    'class': 'subnetwork',
                    'from': [],
                    'subnetwork': {
                      'random': {
                        'class': 'random',
                        'shape': [
                          batch_dim,
                          feat_dim
                        ],
                        'distribution': 'uniform',
                        'minval': 0.0,
                        'maxval': 1.0,
                      },
                      'output': {
                        'class': 'copy',
                        'from': 'random',
                        'out_shape': {batch_dim, feat_dim}
                      }
                    },
                    'out_shape': {batch_dim, feat_dim}
                  },
                  'log': {
                    'class': 'activation',
                    'from': 'random_0',
                    'activation': 'log',
                    'out_shape': {batch_dim, feat_dim}
                  },
                  'negative': {
                    'class': 'activation',
                    'from': 'log',
                    'activation': 'negative',
                    'out_shape': {batch_dim, feat_dim}
                  },
                  'log_0': {
                    'class': 'activation',
                    'from': 'negative',
                    'activation': 'log',
                    'out_shape': {batch_dim, feat_dim}
                  },
                  'negative_0': {
                    'class': 'activation',
                    'from': 'log_0',
                    'activation': 'negative',
                    'out_shape': {batch_dim, feat_dim}
                  },
                  'reduce': {
                    'class': 'reduce',
                    'from': 'random',
                    'mode': 'max',
                    'axis': (batch_dim,),
                    'out_shape': {}
                  },
                  'top_k': {
                    'class': 'top_k',
                    'from': 'negative_0',
                    'axis': feat_dim,
                    'k': 'reduce',
                    'k_dim': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                    'sorted': True,
                    'out_shape': {batch_dim, specaugment_v2_cond_true_feature_masking_top_k_k_dim}
                  },
                  'loop': {
                    'class': 'rec',
                    'from': [],
                    'unit': {
                      'rec_unstack': {
                        'class': 'rec_unstack',
                        'from': 'base:range_in_axis',
                        'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                        'out_shape': {}
                      },
                      'gather': {
                        'class': 'gather',
                        'from': 'base:top_k/indices',
                        'position': 'rec_unstack',
                        'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                        'out_shape': {batch_dim, ImplicitSparseDim(feat_dim)}
                      },
                      '_mask_v2': {
                        'class': 'subnetwork',
                        'from': [],
                        'subnetwork': {
                          'length': {
                            'class': 'length',
                            'from': 'base:prev:_mask_v2',
                            'axis': feat_dim,
                            'out_shape': {}
                          },
                          'random': {
                            'class': 'subnetwork',
                            'from': [],
                            'subnetwork': {
                              'random': {
                                'class': 'random',
                                'shape': (batch_dim,),
                                'distribution': 'uniform',
                                'minval': 1,
                                'maxval': 11,
                                'dtype': 'int32',
                              },
                              'output': {
                                'class': 'copy',
                                'from': 'random',
                                'out_shape': {batch_dim}
                              }
                            },
                            'out_shape': {batch_dim}
                          },
                          'add': {
                            'class': 'combine',
                            'from': ['base:gather', 'random'],
                            'kind': 'add',
                            'out_shape': {batch_dim, ImplicitSparseDim(feat_dim)}
                          },
                          'minimum': {
                            'class': 'combine',
                            'from': ['add', 'length'],
                            'kind': 'minimum',
                            'out_shape': {batch_dim, ImplicitSparseDim(feat_dim)}
                          },
                          'range_in_axis': {
                            'class': 'range_in_axis',
                            'from': 'base:prev:_mask_v2',
                            'axis': feat_dim,
                            'out_shape': {feat_dim}
                          },
                          'greater_equal': {
                            'class': 'compare',
                            'from': ['range_in_axis', 'base:gather'],
                            'kind': 'greater_equal',
                            'allow_broadcast_all_sources': True,
                            'out_shape': {batch_dim, feat_dim}
                          },
                          'less': {
                            'class': 'compare',
                            'from': ['range_in_axis', 'minimum'],
                            'kind': 'less',
                            'allow_broadcast_all_sources': True,
                            'out_shape': {batch_dim, feat_dim}
                          },
                          'logical_and': {
                            'class': 'combine',
                            'from': ['greater_equal', 'less'],
                            'kind': 'logical_and',
                            'out_shape': {batch_dim, feat_dim}
                          },
                          'where': {
                            'class': 'switch',
                            'condition': 'logical_and',
                            'true_from': 0.0,
                            'false_from': 'base:prev:_mask_v2',
                            'out_shape': {batch_dim, time_dim, feat_dim}
                          },
                          'output': {
                            'class': 'copy',
                            'from': 'where',
                            'out_shape': {batch_dim, time_dim, feat_dim}
                          }
                        },
                        'initial_output': 'base:base:base:base:data',
                        'need_last': True,
                        'out_shape': {batch_dim, time_dim, feat_dim}
                      },
                      'output': {
                        'class': 'copy',
                        'from': 'rec_unstack',
                        'out_shape': {}
                      }
                    },
                    'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                    'out_shape': {specaugment_v2_cond_true_feature_masking_top_k_k_dim},
                    'name_scope': ''
                  },
                  'range_in_axis': {
                    'class': 'range_in_axis',
                    'from': 'top_k/indices',
                    'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                    'out_shape': {specaugment_v2_cond_true_feature_masking_top_k_k_dim}
                  },
                  '_mask_v2': {
                    'class': 'rec_last_output',
                    'rec_layer': 'loop',
                    'sub_layer_name': '_mask_v2',
                    'out_shape': {batch_dim, time_dim, feat_dim}
                  },
                  'output': {
                    'class': 'copy',
                    'from': '_mask_v2',
                    'out_shape': {batch_dim, time_dim, feat_dim}
                  }
                },
                'out_shape': {batch_dim, time_dim, feat_dim}
              },
              'output': {
                'class': 'copy',
                'from': 'feature_masking',
                'out_shape': {batch_dim, time_dim, feat_dim}
              }
            }
          },
          'false_layer': {
            'class': 'subnetwork',
            'from': [],
            'subnetwork': {
              'output': {
                'class': 'copy',
                'from': 'base:base:data:data',
                'out_shape': {batch_dim, time_dim, feat_dim}
              }
            }
          },
          'out_shape': {batch_dim, time_dim, feat_dim},
          'name_scope': ''
        },
        "output": {"class": "copy", "from": "cond"},
      },
      'out_shape': {batch_dim, time_dim, feat_dim}
    },
    'output': {
      'class': 'copy',
      'from': 'specaugment_v2',
      'out_shape': {batch_dim, time_dim, feat_dim}
    }
  }

  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    net.initialize_params(session)
    session.run(net.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_specaugment_pure_returnn_reduced():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim, ImplicitSparseDim

  time_dim = SpatialDim('time')
  feat_dim = FeatureDim('feat', 50)

  config = Config(dict(
    extern_data={
      'data': {
        'dim_tags': (batch_dim, time_dim, feat_dim),
        'dtype': 'float32',
        'available_for_inference': True
      }
    },
    debug_runtime_sanity_checks=True,
  ))

  specaugment_v2_cond_true_feature_masking_top_k_k_dim = SpatialDim('feature_masking:num')

  net_dict = {
    'specaugment_v2': {
      'class': 'subnetwork',
      'from': [],
      'subnetwork': {
        'dim_value': {
          'class': 'subnetwork',
          'from': [],
          'subnetwork': {
            'length': {
              'class': 'length',
              'from': 'base:base:data:data',
              'axis': time_dim,
              'out_shape': {batch_dim}
            },
            'reduce': {
              'class': 'reduce',
              'from': 'length',
              'mode': 'max',
              'axis': (batch_dim,),
              'out_shape': {}
            },
            'output': {
              'class': 'copy',
              'from': 'reduce',
              'out_shape': {}
            }
          },
          'out_shape': {}
        },
        'constant': {'class': 'constant', 'value': 2},
        'minimum': {
          'class': 'combine',
          'from': ['constant', 'dim_value'],
          'kind': 'minimum',
          'out_shape': {}
        },
        'constant_0': {'class': 'constant', 'value': 100},
        'floordiv': {
          'class': 'combine',
          'from': ['dim_value', 'constant_0'],
          'kind': 'floordiv',
          'out_shape': {}
        },
        'constant_1': {'class': 'constant', 'value': 2},
        'maximum': {
          'class': 'combine',
          'from': ['floordiv', 'constant_1'],
          'kind': 'maximum',
          'out_shape': {}
        },
        'constant_2': {'class': 'constant', 'value': 4},
        'mul': {
          'class': 'combine',
          'from': ['maximum', 'constant_2'],
          'kind': 'mul',
          'out_shape': {}
        },
        'minimum_0': {
          'class': 'combine',
          'from': ['mul', 'dim_value'],
          'kind': 'minimum',
          'out_shape': {}
        },
        'feature_masking': {
          'class': 'subnetwork',
          'from': [],
          'subnetwork': {
            'random': {
              'class': 'subnetwork',
              'from': [],
              'subnetwork': {
                'random': {
                  'class': 'random',
                  'shape': [
                    batch_dim
                  ],
                  'distribution': 'uniform',
                  'minval': 2,
                  'maxval': 6,
                  'dtype': 'int32',
                },
                'output': {
                  'class': 'copy',
                  'from': 'random',
                  'out_shape': {batch_dim}
                }
              },
              'out_shape': {batch_dim}
            },
            'random_0': {
              'class': 'subnetwork',
              'from': [],
              'subnetwork': {
                'random': {
                  'class': 'random',
                  'shape': [
                    batch_dim,
                    feat_dim
                  ],
                  'distribution': 'uniform',
                  'minval': 0.0,
                  'maxval': 1.0,
                },
                'output': {
                  'class': 'copy',
                  'from': 'random',
                  'out_shape': {batch_dim, feat_dim}
                }
              },
              'out_shape': {batch_dim, feat_dim}
            },
            'log': {
              'class': 'activation',
              'from': 'random_0',
              'activation': 'log',
              'out_shape': {batch_dim, feat_dim}
            },
            'negative': {
              'class': 'activation',
              'from': 'log',
              'activation': 'negative',
              'out_shape': {batch_dim, feat_dim}
            },
            'log_0': {
              'class': 'activation',
              'from': 'negative',
              'activation': 'log',
              'out_shape': {batch_dim, feat_dim}
            },
            'negative_0': {
              'class': 'activation',
              'from': 'log_0',
              'activation': 'negative',
              'out_shape': {batch_dim, feat_dim}
            },
            'reduce': {
              'class': 'reduce',
              'from': 'random',
              'mode': 'max',
              'axis': (batch_dim,),
              'out_shape': {}
            },
            'top_k': {
              'class': 'top_k',
              'from': 'negative_0',
              'axis': feat_dim,
              'k': 'reduce',
              'k_dim': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
              'sorted': True,
              'out_shape': {batch_dim, specaugment_v2_cond_true_feature_masking_top_k_k_dim}
            },
            'loop': {
              'class': 'rec',
              'from': [],
              'unit': {
                'rec_unstack': {
                  'class': 'rec_unstack',
                  'from': 'base:range_in_axis',
                  'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                  'out_shape': {}
                },
                'gather': {
                  'class': 'gather',
                  'from': 'base:top_k/indices',
                  'position': 'rec_unstack',
                  'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                  'out_shape': {batch_dim, ImplicitSparseDim(feat_dim)}
                },
                '_mask_v2': {
                  'class': 'subnetwork',
                  'from': [],
                  'subnetwork': {
                    'length': {
                      'class': 'length',
                      'from': 'base:prev:_mask_v2',
                      'axis': feat_dim,
                      'out_shape': {}
                    },
                    'random': {
                      'class': 'subnetwork',
                      'from': [],
                      'subnetwork': {
                        'random': {
                          'class': 'random',
                          'shape': (batch_dim,),
                          'distribution': 'uniform',
                          'minval': 1,
                          'maxval': 11,
                          'dtype': 'int32',
                        },
                        'output': {
                          'class': 'copy',
                          'from': 'random',
                          'out_shape': {batch_dim}
                        }
                      },
                      'out_shape': {batch_dim}
                    },
                    'add': {
                      'class': 'combine',
                      'from': ['base:gather', 'random'],
                      'kind': 'add',
                      'out_shape': {batch_dim, ImplicitSparseDim(feat_dim)}
                    },
                    'minimum': {
                      'class': 'combine',
                      'from': ['add', 'length'],
                      'kind': 'minimum',
                      'out_shape': {batch_dim, ImplicitSparseDim(feat_dim)}
                    },
                    'range_in_axis': {
                      'class': 'range_in_axis',
                      'from': 'base:prev:_mask_v2',
                      'axis': feat_dim,
                      'out_shape': {feat_dim}
                    },
                    'greater_equal': {
                      'class': 'compare',
                      'from': ['range_in_axis', 'base:gather'],
                      'kind': 'greater_equal',
                      'allow_broadcast_all_sources': True,
                      'out_shape': {batch_dim, feat_dim}
                    },
                    'less': {
                      'class': 'compare',
                      'from': ['range_in_axis', 'minimum'],
                      'kind': 'less',
                      'allow_broadcast_all_sources': True,
                      'out_shape': {batch_dim, feat_dim}
                    },
                    'logical_and': {
                      'class': 'combine',
                      'from': ['greater_equal', 'less'],
                      'kind': 'logical_and',
                      'out_shape': {batch_dim, feat_dim}
                    },
                    'where': {
                      'class': 'switch',
                      'condition': 'logical_and',
                      'true_from': 0.0,
                      'false_from': 'base:prev:_mask_v2',
                      'out_shape': {batch_dim, time_dim, feat_dim}
                    },
                    'output': {
                      'class': 'copy',
                      'from': 'where',
                      'out_shape': {batch_dim, time_dim, feat_dim}
                    }
                  },
                  'initial_output': 'base:base:base:data',
                  'need_last': True,
                  'out_shape': {batch_dim, time_dim, feat_dim}
                },
                'output': {
                  'class': 'copy',
                  'from': 'rec_unstack',
                  'out_shape': {}
                }
              },
              'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
              'out_shape': {specaugment_v2_cond_true_feature_masking_top_k_k_dim},
              'name_scope': ''
            },
            'range_in_axis': {
              'class': 'range_in_axis',
              'from': 'top_k/indices',
              'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
              'out_shape': {specaugment_v2_cond_true_feature_masking_top_k_k_dim}
            },
            '_mask_v2': {
              'class': 'rec_last_output',
              'rec_layer': 'loop',
              'sub_layer_name': '_mask_v2',
              'out_shape': {batch_dim, time_dim, feat_dim}
            },
            'output': {
              'class': 'copy',
              'from': '_mask_v2',
              'out_shape': {batch_dim, time_dim, feat_dim}
            }
          },
          'out_shape': {batch_dim, time_dim, feat_dim}
        },
        'output': {
          'class': 'copy',
          'from': 'feature_masking',
          'out_shape': {batch_dim, time_dim, feat_dim}
        }
      },
      'out_shape': {batch_dim, time_dim, feat_dim}
    },
    'output': {
      'class': 'copy',
      'from': 'specaugment_v2',
      'out_shape': {batch_dim, time_dim, feat_dim}
    }
  }

  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    net.initialize_params(session)
    session.run(net.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_specaugment_pure_returnn():
  from returnn.tf.util.data import (
    batch_dim, SpatialDim, FeatureDim, ImplicitDynSizeDim, ImplicitSparseDim)

  time_dim = SpatialDim('time')
  feat_dim = FeatureDim('feat', 50)

  config = Config(dict(
    extern_data={
      'data': {
        'dim_tags': (batch_dim, time_dim, feat_dim),
        'dtype': 'float32',
        'available_for_inference': True
      }
    },
    debug_runtime_sanity_checks=True,
  ))

  specaugment_v2_cond_true_time_masking_top_k_k_dim = SpatialDim('specaugment_v2/cond/true/time_masking:top_k:k_dim')
  specaugment_v2_cond_true_feature_masking_top_k_k_dim = SpatialDim(
    'specaugment_v2/cond/true/feature_masking:top_k:k_dim')
  random_state_dim = FeatureDim('random-state', 3)

  net_dict = {
    'specaugment_v2': {
      'class': 'subnetwork',
      'from': [],
      'subnetwork': {
        'ones': {'class': 'constant', 'value': True, 'shape': (), 'dtype': 'bool'},
        'cond': {
          'class': 'cond',
          'from': [],
          'condition': 'ones',
          'true_layer': {
            'class': 'subnetwork',
            'from': [],
            'subnetwork': {
              'dim_value': {
                'class': 'subnetwork',
                'from': [],
                'subnetwork': {
                  'length': {
                    'class': 'length',
                    'from': 'base:base:base:data:data',
                    'axis': time_dim,
                    'out_shape': {batch_dim}
                  },
                  'reduce': {
                    'class': 'reduce',
                    'from': 'length',
                    'mode': 'max',
                    'axis': (batch_dim,),
                    'out_shape': {}
                  },
                  'output': {
                    'class': 'copy',
                    'from': 'reduce',
                    'out_shape': {}
                  }
                },
                'out_shape': {}
              },
              'constant': {'class': 'constant', 'value': 2},
              'minimum': {
                'class': 'combine',
                'from': ['constant', 'dim_value'],
                'kind': 'minimum',
                'out_shape': {}
              },
              'constant_0': {'class': 'constant', 'value': 100},
              'floordiv': {
                'class': 'combine',
                'from': ['dim_value', 'constant_0'],
                'kind': 'floordiv',
                'out_shape': {}
              },
              'constant_1': {'class': 'constant', 'value': 2},
              'maximum': {
                'class': 'combine',
                'from': ['floordiv', 'constant_1'],
                'kind': 'maximum',
                'out_shape': {}
              },
              'constant_2': {'class': 'constant', 'value': 4},
              'mul': {
                'class': 'combine',
                'from': ['maximum', 'constant_2'],
                'kind': 'mul',
                'out_shape': {}
              },
              'minimum_0': {
                'class': 'combine',
                'from': ['mul', 'dim_value'],
                'kind': 'minimum',
                'out_shape': {}
              },
              'time_masking': {
                'class': 'subnetwork',
                'from': [],
                'subnetwork': {
                  'constant': {'class': 'constant', 'value': 1},
                  'add': {
                    'class': 'combine',
                    'from': ['base:minimum_0', 'constant'],
                    'kind': 'add',
                    'out_shape': {}
                  },
                  'random': {
                    'class': 'subnetwork',
                    'from': [],
                    'subnetwork': {
                      'random': {
                        'class': 'random',
                        'shape': [
                          batch_dim
                        ],
                        'distribution': 'uniform',
                        'minval': 'base:base:minimum',
                        'maxval': 'base:add',
                        'dtype': 'int32',
                        'explicit_state': 'base:base:base:random_state_var0_1',
                        'auto_update_state': True
                      },
                      'output': {
                        'class': 'copy',
                        'from': 'random',
                        'out_shape': {batch_dim}
                      }
                    },
                    'out_shape': {batch_dim}
                  },
                  'random_0': {
                    'class': 'subnetwork',
                    'from': [],
                    'subnetwork': {
                      'random': {
                        'class': 'random',
                        'shape': [
                          batch_dim,
                          time_dim
                        ],
                        'distribution': 'uniform',
                        'minval': 0.0,
                        'maxval': 1.0,
                        'explicit_state': 'base:base:base:random_0_state_var0_0',
                        'auto_update_state': True
                      },
                      'output': {
                        'class': 'copy',
                        'from': 'random',
                        'out_shape': {batch_dim, time_dim}
                      }
                    },
                    'out_shape': {batch_dim, time_dim}
                  },
                  'log': {
                    'class': 'activation',
                    'from': 'random_0',
                    'activation': 'log',
                    'out_shape': {batch_dim, time_dim}
                  },
                  'negative': {
                    'class': 'activation',
                    'from': 'log',
                    'activation': 'negative',
                    'out_shape': {batch_dim, time_dim}
                  },
                  'log_0': {
                    'class': 'activation',
                    'from': 'negative',
                    'activation': 'log',
                    'out_shape': {batch_dim, time_dim}
                  },
                  'negative_0': {
                    'class': 'activation',
                    'from': 'log_0',
                    'activation': 'negative',
                    'out_shape': {batch_dim, time_dim}
                  },
                  'reduce': {
                    'class': 'reduce',
                    'from': 'random',
                    'mode': 'max',
                    'axis': (batch_dim,),
                    'out_shape': {}
                  },
                  'top_k': {
                    'class': 'top_k',
                    'from': 'negative_0',
                    'axis': time_dim,
                    'k': 'reduce',
                    'k_dim': specaugment_v2_cond_true_time_masking_top_k_k_dim,
                    'sorted': True,
                    'out_shape': {batch_dim, specaugment_v2_cond_true_time_masking_top_k_k_dim}
                  },
                  'loop': {
                    'class': 'rec',
                    'from': [],
                    'unit': {
                      'rec_unstack': {
                        'class': 'rec_unstack',
                        'from': 'base:range_in_axis',
                        'axis': specaugment_v2_cond_true_time_masking_top_k_k_dim,
                        'out_shape': {}
                      },
                      'gather': {
                        'class': 'gather',
                        'from': 'base:top_k/indices',
                        'position': 'rec_unstack',
                        'axis': specaugment_v2_cond_true_time_masking_top_k_k_dim,
                        'out_shape': {batch_dim, ImplicitSparseDim(time_dim)}
                      },
                      '_mask_v2': {
                        'class': 'subnetwork',
                        'from': [],
                        'subnetwork': {
                          'length': {
                            'class': 'length',
                            'from': 'base:prev:_mask_v2',
                            'axis': time_dim,
                            'out_shape': {batch_dim}
                          },
                          'random': {
                            'class': 'subnetwork',
                            'from': [],
                            'subnetwork': {
                              'random': {
                                'class': 'random',
                                'shape': (batch_dim,),
                                'distribution': 'uniform',
                                'minval': 1,
                                'maxval': 21,
                                'dtype': 'int32',
                                'explicit_state': 'base:base:base:base:base:random_state_var0_2',
                                'auto_update_state': True
                              },
                              'output': {
                                'class': 'copy',
                                'from': 'random',
                                'out_shape': {batch_dim}
                              }
                            },
                            'out_shape': {batch_dim}
                          },
                          'add': {
                            'class': 'combine',
                            'from': ['base:gather', 'random'],
                            'kind': 'add',
                            'out_shape': {batch_dim, ImplicitSparseDim(time_dim)}
                          },
                          'minimum': {
                            'class': 'combine',
                            'from': ['add', 'length'],
                            'kind': 'minimum',
                            'out_shape': {batch_dim, ImplicitSparseDim(time_dim)}
                          },
                          'range_in_axis': {
                            'class': 'range_in_axis',
                            'from': 'base:prev:_mask_v2',
                            'axis': time_dim,
                            'out_shape': {ImplicitDynSizeDim(batch_dim), time_dim}
                          },
                          'greater_equal': {
                            'class': 'compare',
                            'from': ['range_in_axis', 'base:gather'],
                            'kind': 'greater_equal',
                            'allow_broadcast_all_sources': True,
                            'out_shape': {batch_dim, time_dim}
                          },
                          'less': {
                            'class': 'compare',
                            'from': ['range_in_axis', 'minimum'],
                            'kind': 'less',
                            'allow_broadcast_all_sources': True,
                            'out_shape': {batch_dim, time_dim}
                          },
                          'logical_and': {
                            'class': 'combine',
                            'from': ['greater_equal', 'less'],
                            'kind': 'logical_and',
                            'out_shape': {batch_dim, time_dim}
                          },
                          'where': {
                            'class': 'switch',
                            'condition': 'logical_and',
                            'true_from': 0.0,
                            'false_from': 'base:prev:_mask_v2',
                            'out_shape': {batch_dim, time_dim, feat_dim}
                          },
                          'output': {
                            'class': 'copy',
                            'from': 'where',
                            'out_shape': {batch_dim, time_dim, feat_dim}
                          }
                        },
                        'initial_output': 'base:base:base:base:data:data',
                        'need_last': True,
                        'out_shape': {batch_dim, time_dim, feat_dim}
                      },
                      'output': {
                        'class': 'copy',
                        'from': 'rec_unstack',
                        'out_shape': {}
                      }
                    },
                    'axis': specaugment_v2_cond_true_time_masking_top_k_k_dim,
                    'out_shape': {specaugment_v2_cond_true_time_masking_top_k_k_dim},
                    'name_scope': ''
                  },
                  'range_in_axis': {
                    'class': 'range_in_axis',
                    'from': 'top_k/indices',
                    'axis': specaugment_v2_cond_true_time_masking_top_k_k_dim,
                    'out_shape': {specaugment_v2_cond_true_time_masking_top_k_k_dim}
                  },
                  '_mask_v2': {
                    'class': 'rec_last_output',
                    'rec_layer': 'loop',
                    'sub_layer_name': '_mask_v2',
                    'out_shape': {batch_dim, time_dim, feat_dim}
                  },
                  'output': {
                    'class': 'copy',
                    'from': '_mask_v2',
                    'out_shape': {batch_dim, time_dim, feat_dim}
                  }
                },
                'out_shape': {batch_dim, time_dim, feat_dim}
              },
              'feature_masking': {
                'class': 'subnetwork',
                'from': [],
                'subnetwork': {
                  'random': {
                    'class': 'subnetwork',
                    'from': [],
                    'subnetwork': {
                      'random': {
                        'class': 'random',
                        'shape': [
                          batch_dim
                        ],
                        'distribution': 'uniform',
                        'minval': 2,
                        'maxval': 6,
                        'dtype': 'int32',
                        'explicit_state': 'base:base:base:random_state_var0',
                        'auto_update_state': True
                      },
                      'output': {
                        'class': 'copy',
                        'from': 'random',
                        'out_shape': {batch_dim}
                      }
                    },
                    'out_shape': {batch_dim}
                  },
                  'random_0': {
                    'class': 'subnetwork',
                    'from': [],
                    'subnetwork': {
                      'random': {
                        'class': 'random',
                        'shape': [
                          batch_dim,
                          feat_dim
                        ],
                        'distribution': 'uniform',
                        'minval': 0.0,
                        'maxval': 1.0,
                        'explicit_state': 'base:base:base:random_0_state_var0',
                        'auto_update_state': True
                      },
                      'output': {
                        'class': 'copy',
                        'from': 'random',
                        'out_shape': {batch_dim, feat_dim}
                      }
                    },
                    'out_shape': {batch_dim, feat_dim}
                  },
                  'log': {
                    'class': 'activation',
                    'from': 'random_0',
                    'activation': 'log',
                    'out_shape': {batch_dim, feat_dim}
                  },
                  'negative': {
                    'class': 'activation',
                    'from': 'log',
                    'activation': 'negative',
                    'out_shape': {batch_dim, feat_dim}
                  },
                  'log_0': {
                    'class': 'activation',
                    'from': 'negative',
                    'activation': 'log',
                    'out_shape': {batch_dim, feat_dim}
                  },
                  'negative_0': {
                    'class': 'activation',
                    'from': 'log_0',
                    'activation': 'negative',
                    'out_shape': {batch_dim, feat_dim}
                  },
                  'reduce': {
                    'class': 'reduce',
                    'from': 'random',
                    'mode': 'max',
                    'axis': (batch_dim,),
                    'out_shape': {}
                  },
                  'top_k': {
                    'class': 'top_k',
                    'from': 'negative_0',
                    'axis': feat_dim,
                    'k': 'reduce',
                    'k_dim': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                    'sorted': True,
                    'out_shape': {batch_dim, specaugment_v2_cond_true_feature_masking_top_k_k_dim}
                  },
                  'loop': {
                    'class': 'rec',
                    'from': [],
                    'unit': {
                      'rec_unstack': {
                        'class': 'rec_unstack',
                        'from': 'base:range_in_axis',
                        'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                        'out_shape': {}
                      },
                      'gather': {
                        'class': 'gather',
                        'from': 'base:top_k/indices',
                        'position': 'rec_unstack',
                        'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                        'out_shape': {batch_dim, ImplicitSparseDim(feat_dim)}
                      },
                      '_mask_v2': {
                        'class': 'subnetwork',
                        'from': [],
                        'subnetwork': {
                          'length': {
                            'class': 'length',
                            'from': 'base:prev:_mask_v2',
                            'axis': feat_dim,
                            'out_shape': {}
                          },
                          'random': {
                            'class': 'subnetwork',
                            'from': [],
                            'subnetwork': {
                              'random': {
                                'class': 'random',
                                'shape': (batch_dim,),
                                'distribution': 'uniform',
                                'minval': 1,
                                'maxval': 11,
                                'dtype': 'int32',
                                'explicit_state': 'base:base:base:base:base:random_state_var0_0',
                                'auto_update_state': True
                              },
                              'output': {
                                'class': 'copy',
                                'from': 'random',
                                'out_shape': {batch_dim}
                              }
                            },
                            'out_shape': {batch_dim}
                          },
                          'add': {
                            'class': 'combine',
                            'from': ['base:gather', 'random'],
                            'kind': 'add',
                            'out_shape': {batch_dim, ImplicitSparseDim(feat_dim)}
                          },
                          'minimum': {
                            'class': 'combine',
                            'from': ['add', 'length'],
                            'kind': 'minimum',
                            'out_shape': {batch_dim, ImplicitSparseDim(feat_dim)}
                          },
                          'range_in_axis': {
                            'class': 'range_in_axis',
                            'from': 'base:prev:_mask_v2',
                            'axis': feat_dim,
                            'out_shape': {feat_dim}
                          },
                          'greater_equal': {
                            'class': 'compare',
                            'from': ['range_in_axis', 'base:gather'],
                            'kind': 'greater_equal',
                            'allow_broadcast_all_sources': True,
                            'out_shape': {batch_dim, feat_dim}
                          },
                          'less': {
                            'class': 'compare',
                            'from': ['range_in_axis', 'minimum'],
                            'kind': 'less',
                            'allow_broadcast_all_sources': True,
                            'out_shape': {batch_dim, feat_dim}
                          },
                          'logical_and': {
                            'class': 'combine',
                            'from': ['greater_equal', 'less'],
                            'kind': 'logical_and',
                            'out_shape': {batch_dim, feat_dim}
                          },
                          'where': {
                            'class': 'switch',
                            'condition': 'logical_and',
                            'true_from': 0.0,
                            'false_from': 'base:prev:_mask_v2',
                            'out_shape': {batch_dim, time_dim, feat_dim}
                          },
                          'output': {
                            'class': 'copy',
                            'from': 'where',
                            'out_shape': {batch_dim, time_dim, feat_dim}
                          }
                        },
                        'initial_output': 'base:base:time_masking',
                        'need_last': True,
                        'out_shape': {batch_dim, time_dim, feat_dim}
                      },
                      'output': {
                        'class': 'copy',
                        'from': 'rec_unstack',
                        'out_shape': {}
                      }
                    },
                    'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                    'out_shape': {specaugment_v2_cond_true_feature_masking_top_k_k_dim},
                    'name_scope': ''
                  },
                  'range_in_axis': {
                    'class': 'range_in_axis',
                    'from': 'top_k/indices',
                    'axis': specaugment_v2_cond_true_feature_masking_top_k_k_dim,
                    'out_shape': {specaugment_v2_cond_true_feature_masking_top_k_k_dim}
                  },
                  '_mask_v2': {
                    'class': 'rec_last_output',
                    'rec_layer': 'loop',
                    'sub_layer_name': '_mask_v2',
                    'out_shape': {batch_dim, time_dim, feat_dim}
                  },
                  'output': {
                    'class': 'copy',
                    'from': '_mask_v2',
                    'out_shape': {batch_dim, time_dim, feat_dim}
                  }
                },
                'out_shape': {batch_dim, time_dim, feat_dim}
              },
              'output': {
                'class': 'copy',
                'from': 'feature_masking',
                'out_shape': {batch_dim, time_dim, feat_dim}
              }
            }
          },
          'false_layer': {
            'class': 'subnetwork',
            'from': [],
            'subnetwork': {
              'output': {
                'class': 'copy',
                'from': 'base:base:data:data',
                'out_shape': {batch_dim, time_dim, feat_dim}
              }
            }
          },
          'out_shape': {batch_dim, time_dim, feat_dim},
          'name_scope': ''
        },
        'random_state_init': {
          'class': 'random_state_init',
          'out_dim': random_state_dim,
          'out_shape': {random_state_dim}
        },
        'random_state_init_0': {
          'class': 'random_state_init',
          'out_dim': random_state_dim,
          'out_shape': {random_state_dim}
        },
        'random_state_init_1': {
          'class': 'random_state_init',
          'out_dim': random_state_dim,
          'out_shape': {random_state_dim}
        },
        'random_state_init_2': {
          'class': 'random_state_init',
          'out_dim': random_state_dim,
          'out_shape': {random_state_dim}
        },
        'random_state_init_3': {
          'class': 'random_state_init',
          'out_dim': random_state_dim,
          'out_shape': {random_state_dim}
        },
        'random_state_init_4': {
          'class': 'random_state_init',
          'out_dim': random_state_dim,
          'out_shape': {random_state_dim}
        },
        'output': {
          'class': 'copy',
          'from': 'cond',
          'out_shape': {batch_dim, time_dim, feat_dim}
        },
        'random_state_var0': {
          'class': 'variable',
          'shape': [
            random_state_dim
          ],
          'param_name': 'param',
          'dtype': 'int64',
          'init_by_layer': 'random_state_init_2'
        },
        'random_state_var0_0': {
          'class': 'variable',
          'shape': [
            random_state_dim
          ],
          'param_name': 'param',
          'dtype': 'int64',
          'init_by_layer': 'random_state_init_4'
        },
        'random_0_state_var0': {
          'class': 'variable',
          'shape': [
            random_state_dim
          ],
          'param_name': 'param',
          'dtype': 'int64',
          'init_by_layer': 'random_state_init_3'
        },
        'random_state_var0_1': {
          'class': 'variable',
          'shape': [
            random_state_dim
          ],
          'param_name': 'param',
          'dtype': 'int64',
          'init_by_layer': 'random_state_init'
        },
        'random_state_var0_2': {
          'class': 'variable',
          'shape': [
            random_state_dim
          ],
          'param_name': 'param',
          'dtype': 'int64',
          'init_by_layer': 'random_state_init_1'
        },
        'random_0_state_var0_0': {
          'class': 'variable',
          'shape': [
            random_state_dim
          ],
          'param_name': 'param',
          'dtype': 'int64',
          'init_by_layer': 'random_state_init_0'
        }
      },
      'out_shape': {batch_dim, time_dim, feat_dim}
    },
    'output': {
      'class': 'copy',
      'from': 'specaugment_v2',
      'out_shape': {batch_dim, time_dim, feat_dim}
    }
  }

  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    net.initialize_params(session)
    session.run(net.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_SearchSortedLayer():
  n_batch, n_time, n_in, n_out = 2, 10, 3, 5
  random = numpy.random.RandomState(seed=1)
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData({"data": {"shape": (None, 3)}}))
    sorted_layer = InternalLayer(
      name="sorted_sequence", network=net,
      output=Data(name="sorted_sequence", shape=(None, n_in)))  # [B,T,F]
    sorted = numpy.sort(random.uniform(0.0, 10.0, size=(n_batch, n_time, n_in)), axis=1)
    sorted_lens = [10, 7]
    sorted_layer.output.placeholder = tf.constant(sorted, dtype="float32")
    sorted_layer.output.size_placeholder = {0: tf.constant(sorted_lens, dtype="int32")}  # [B]
    values_layer = InternalLayer(
      name="values", network=net,
      output=Data(
        **{"name": "values", "shape": (n_in, n_out), "time_dim_axis": None, "feature_dim_axis": 1}))  # [B,F,F']
    values = random.uniform(0.0, 10.0, size=(n_batch, n_in, n_out))
    values_layer.output.placeholder = tf.constant(values, dtype="float32")

    for side in ["left", "right"]:
      print("Testing side=%r" % side)
      opts = {
        "network": net, "name": 'search_sorted_test', "sources": [], "sorted_sequence": sorted_layer,
        "values": values_layer, "axis": "T", "side": side}
      out_data = SearchSortedLayer.get_out_data_from_opts(**opts)
      out_data.sanity_check(ignore_placeholder=True)
      search_layer = SearchSortedLayer(output=out_data, **opts)  # should be [B,F]
      out_data = search_layer.output
      out_data.sanity_check()
      print(search_layer.output)

      output = session.run(search_layer.output.placeholder, feed_dict=make_feed_dict(net.extern_data, n_batch=n_batch))
      assert output.dtype == "int32"
      assert out_data.batch_shape == (None, n_in, n_out)
      assert out_data.batch_dim_axis == 0
      assert out_data.time_dim_axis is None
      assert out_data.feature_dim_axis == 1
      for b, t_max in enumerate(sorted_lens):
        for f in range(n_in):
          expected = numpy.searchsorted(a=sorted[b, :t_max, f], v=values[b, f, :], side=side)
          actual = output[b, f, :]
          numpy.testing.assert_equal(expected, actual)


def test_CondLayer_subnetwork_train():
  n_batch, n_time, n_in, n_out = 3, 7, 11, 13
  config = Config({
    "extern_data": {
      "data": {"dim": n_in},
      "classes": {"dim": n_out, "sparse": True},
    },
    "debug_print_layer_output_template": True,
  })
  rnd = numpy.random.RandomState(42)
  with make_scope() as session:
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "src": {"class": "linear", "activation": "tanh", "n_out": 10, "from": "data"},
      "cond": {
        "class": "cond", "from": [],
        "condition": {
          "class": "eval", "from": [], "out_type": {"batch_dim_axis": None, "shape": (), "dtype": "bool"},
          "eval": "tf.equal(self.network.global_train_step % 2, 0)"
          },
        "true_layer": {
          "class": "subnetwork", "from": "src", "subnetwork": {
            "lin": {"class": "linear", "activation": "tanh", "n_out": 10, "from": "data"},
            "res": {"class": "combine", "kind": "add", "from": ["data", "lin"]},
            "output": {"class": "print", "from": "res", "extra_print_args": ["true_layer"], "summarize": 1}
          }},
        "false_layer": {"class": "copy", "from": "src"}
        },
      "output": {"class": "softmax", "from": "cond", "loss": "ce", "target": "classes"}})
    net.print_network_info()
    cond_layer = net.get_layer("cond")
    assert isinstance(cond_layer, CondLayer)
    assert not tf_util.has_control_flow_context(cond_layer.output.placeholder)
    cond_true_layer = cond_layer.true_layer
    assert isinstance(cond_true_layer, SubnetworkLayer)
    # Check whether the execution of the true branch is actually conditionally.
    assert tf_util.has_control_flow_context(cond_true_layer.output.placeholder)
    trainable_vars = net.get_trainable_params()
    print("Trainable vars:")
    pprint(trainable_vars)
    cond_var = net.layers["cond"].params["lin/W"]
    assert cond_var in trainable_vars
    from returnn.tf.updater import Updater
    updater = Updater(config=config, network=net, initial_learning_rate=0.1)
    updater.set_trainable_vars(trainable_vars)
    updater.init_optimizer_vars(session)
    updater.set_learning_rate(value=updater.initial_learning_rate, session=session)
    net.initialize_params(session)
    in_v = rnd.normal(size=(n_batch, n_time, n_in)).astype("float32")
    targets_v = rnd.randint(0, n_out, size=(n_batch, n_time)).astype("int32")
    seq_lens_v = numpy.array([n_time, n_time - 1, n_time - 2])
    assert len(seq_lens_v) == n_batch
    feed_dict = {
      net.extern_data.data["data"].placeholder: in_v,
      net.extern_data.data["data"].size_placeholder[0]: seq_lens_v,
      net.extern_data.data["classes"].placeholder: targets_v,
      net.extern_data.data["classes"].size_placeholder[0]: seq_lens_v,
    }
    fetches = net.get_fetches_dict(with_summary=True, with_size=True)
    fetches["optim_op"] = updater.get_optim_op()
    try:
      loss = None
      initial_loss = float("inf")
      for i in range(10):
        step = session.run(net.global_train_step)
        print("step: %i" % step)
        assert i == step
        old_var_value = session.run(cond_var)
        result = session.run(feed_dict=feed_dict, fetches=fetches)
        loss = result["loss"]
        print("loss:", loss)
        if i == 0:
          initial_loss = loss
        new_var_value = session.run(cond_var)
        var_changed = (old_var_value != new_var_value).any()
        print("var changed:", var_changed)
        if i % 2 == 0:  # See cond layer, condition. Use true_layer every second iteration, starting with 0.
          # We used true_layer, thus the params should have been updated.
          assert var_changed
        else:
          # We did not use true_layer, thus the params should not have been updated.
          assert not var_changed
      assert loss is not None and loss < initial_loss and numpy.isfinite(initial_loss)
    except tf.errors.OpError as exc:
      print("TF exception:", type(exc).__name__, ":", exc)
      from returnn.tf.network import help_on_tf_exception
      help_on_tf_exception(session=session, exception=exc, fetches=fetches, feed_dict=feed_dict)
      raise


def test_CondLayer_subnet_template_construct():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim

  time_dim = SpatialDim('time')
  feat_dim = FeatureDim('feat', 5)

  config = Config(dict(
    extern_data={
      'data': {
        'dim_tags': (batch_dim, time_dim, feat_dim),
        'dtype': 'float32',
        'available_for_inference': True
      }
    },
    debug_runtime_sanity_checks=True,
    debug_print_layer_output_shape=True,
    debug_print_layer_output=True,
  ))

  class _EvalFuncLocals:
    graph_call_count = 0
    session_call_count = 0

  def _py_func(x):
    _EvalFuncLocals.session_call_count += 1
    return x

  def _eval_func(source, **_kwargs):
    _EvalFuncLocals.graph_call_count += 1
    x = source(0)
    y, = tf_compat.v1.py_func(_py_func, [x], [x.dtype], stateful=True)
    y.set_shape(x.get_shape())
    return y

  net_dict = {
    'specaugment_v2': {
      'class': 'subnetwork',
      'from': [],
      'subnetwork': {
        'train_flag': {'class': 'train_flag'},
        'cond': {
          'class': 'cond',
          'from': [],
          'condition': 'train_flag',
          'true_layer': {
            'class': 'subnetwork',
            'from': [],
            'subnetwork': {
              'scores': {
                "class": "reduce", "from": "base:base:data", "mode": "sum",
                "axis": [batch_dim, time_dim, feat_dim],
                "out_shape": {}
              },
              'output': {
                'class': 'eval', "from": "scores", "eval": _eval_func
              },
            }
          },
          'false_layer': {
            'class': 'subnetwork',
            'from': [],
            'subnetwork': {
              'output': {
                'class': 'constant', "value": -1., 'out_shape': {}
              }
            }
          },
          'name_scope': ''
        },
        "output": {"class": "copy", "from": "cond"},
      },
    },
    'output': {
      'class': 'copy',
      'from': 'specaugment_v2',
    }
  }

  with make_scope() as session:
    train_flag = tf_util.get_global_train_flag_placeholder()
    net = TFNetwork(config=config, train_flag=train_flag)
    net.construct_from_dict(net_dict)
    assert _EvalFuncLocals.graph_call_count == 1  # if more often, set a breakpoint above
    net.initialize_params(session)
    feed_dict = make_feed_dict(net.extern_data)
    feed_dict[train_flag] = False
    res_eval = session.run(net.get_default_output_layer().output.placeholder, feed_dict=feed_dict)
    print("eval:", res_eval)
    assert _EvalFuncLocals.session_call_count == 0
    feed_dict[train_flag] = True
    res_train = session.run(net.get_default_output_layer().output.placeholder, feed_dict=feed_dict)
    print("train:", res_train)
    assert _EvalFuncLocals.session_call_count == 1


def test_CondLayer_data_access():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim

  time_dim = SpatialDim('time')
  input_dim = FeatureDim('input', 13)

  config = Config(dict(extern_data={
    'data': {'dim_tags': (batch_dim, time_dim, input_dim)}
  }))
  net_dict = {
    'output': {
      'class': 'copy',
      'from': 'cond',
      'out_shape': {batch_dim, time_dim, input_dim}
    },
    'length': {
      'class': 'length',
      'from': 'data:data',
      'axis': batch_dim,
      'out_shape': {}
    },
    'mod': {
      'class': 'eval',
      'from': 'length',
      'eval': 'source(0) % 2',
      'out_shape': {}
    },
    'eq': {
      'class': 'compare',
      'from': 'mod',
      'kind': 'equal',
      'value': 0,
      'out_shape': {}
    },
    'cond': {
      'class': 'cond',
      'from': [],
      'condition': 'eq',
      'true_layer': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'const': {
            "class": "constant", "value": 1.,
            'shape': (batch_dim, time_dim, input_dim),
            'shape_deps': ['base:data:data']
          },
          'output': {
            'class': 'combine',
            'from': ['base:data:data', 'const'],
            'kind': 'add',
            'out_shape': {batch_dim, time_dim, input_dim}
          },
        }
      },
      'false_layer': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'output': {
            'class': 'copy',
            'from': 'base:data:data',
            'out_shape': {batch_dim, time_dim, input_dim}
          }
        }
      },
      'out_shape': {batch_dim, time_dim, input_dim},
      'name_scope': ''
    },
  }
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    network.initialize_params(session)
    session.run(network.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(network.extern_data))


def test_CondLayer_mult_dyn_axes():
  # https://github.com/rwth-i6/returnn/issues/1207
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim

  time_dim = SpatialDim('time')
  input_dim = FeatureDim('input', 13)
  extra_time_dim = SpatialDim('extra_time')

  config = Config(dict(extern_data={
    'data': {'dim_tags': (batch_dim, time_dim, input_dim)}
  }))
  net_dict = {
    'cond': {
      'class': 'eval',
      'from': 'data',
      'eval': 'tf.equal(tf.shape(source(0, auto_convert=False))[0] % 2, 0)',
      'out_type': {'shape': (), 'dtype': 'bool', 'batch_dim_axis': None, 'time_dim_axis': None},
    },
    'new_dim': {
      'class': 'reinterpret_data',
      'from': 'data',
      'set_dim_tags': {time_dim: extra_time_dim},
    },
    'pool': {
      'class': 'pool',
      'from': 'new_dim',
      'mode': 'max',
      'padding': 'same',
      'pool_size': [2],
    },
    'combine': {
      "class": "combine",
      "from": ["data", "pool"],
      "kind": "mul",
      "out_shape": {batch_dim, time_dim, extra_time_dim.ceildiv_right(2), input_dim},
    },
    'conv': {
      'class': 'conv',
      'from': 'combine',
      'filter_size': [3, 3],
      'strides': 2,
      'padding': 'same',
      'in_spatial_dims': (time_dim, extra_time_dim.ceildiv_right(2)),
      'n_out': 7,
    },
    'output': {
      'class': 'cond',
      'from': [],
      'condition': 'cond',
      'true_layer': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'output': {
            'class': 'eval',
            'from': 'base:conv',
            'eval': 'source(0) * 1.5',
          },
        }
      },
      'false_layer': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'output': {
            'class': 'copy',
            'from': 'base:conv',
          },
        }
      },
    },
  }
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    network.initialize_params(session)
    in_ = network.extern_data.get_default_input_data()
    out = network.get_default_output_layer().output
    t1, _, t2, t3 = session.run(
      (in_.get_time_dim_tag().dyn_size_ext.placeholder,
       out.placeholder,
       out.get_dyn_size_tags()[0].dyn_size_ext.placeholder,
       out.get_dyn_size_tags()[1].dyn_size_ext.placeholder),
      feed_dict=make_feed_dict(network.extern_data))

    def _ceildiv(a, b):
      return -(a // -b)

    numpy.testing.assert_array_equal(t2, _ceildiv(t1, 2))
    numpy.testing.assert_array_equal(t3, _ceildiv(t1, 4))


def test_CondLayer_variational_weight_noise():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim

  time_dim = SpatialDim('time')
  in_dim = FeatureDim('in', 3)

  config = Config(dict(extern_data={
    'data': {
      'dim_tags': (batch_dim, time_dim, in_dim),
    }}))

  out_dim = FeatureDim('out', 5)
  random_state_dim = FeatureDim('random-state', 3)

  net_dict = {
    'output': {
      'class': 'copy',
      'from': 'add',
      'out_shape': {batch_dim, time_dim, out_dim}
    },
    'bias': {
      'class': 'variable',
      'shape': [out_dim],
      'param_name': 'param',
      'init': 0.0
    },
    'weight_raw': {
      'class': 'variable',
      'shape': [in_dim, out_dim],
      'param_name': 'param',
    },
    'train_flag': {
      'class': 'train_flag',
      'out_shape': {}
    },
    'weight': {
      'class': 'cond',
      'from': [],
      'condition': 'train_flag',
      'true_layer': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'random': {
            'class': 'random',
            'shape': (in_dim, out_dim),
            'distribution': 'normal',
            'mean': 0.0,
            'stddev': 0.075,
            'dtype': 'float32',
          },
          'output': {
            'class': 'combine',
            'from': ['base:weight_raw', 'random'],
            'kind': 'add',
            'out_shape': {in_dim, out_dim}
          },
        }
      },
      'false_layer': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'output': {
            'class': 'copy',
            'from': 'base:weight_raw',
            'out_shape': {in_dim, out_dim}
          }
        }
      },
      'out_shape': {in_dim, out_dim},
      'name_scope': ''
    },
    'dot': {
      'class': 'dot',
      'from': ['data:data', 'weight'],
      'reduce': in_dim,
      'out_shape': {batch_dim, time_dim, out_dim}
    },
    'add': {
      'class': 'combine',
      'from': ['dot', 'bias'],
      'kind': 'add',
      'out_shape': {batch_dim, time_dim, out_dim}
    },
  }

  print("* Test with dynamic train flag")
  with make_scope() as session:
    train_flag = tf_compat.v1.placeholder(tf.bool, shape=(), name="train_flag")
    network = TFNetwork(config=config)
    network.construct_from_dict(net_dict)
    network.initialize_params(session)
    feed_dict = make_feed_dict(network.extern_data)
    feed_dict[train_flag] = True
    session.run(network.get_default_output_layer().output.placeholder, feed_dict=feed_dict)
    feed_dict[train_flag] = False
    session.run(network.get_default_output_layer().output.placeholder, feed_dict=feed_dict)

  for train_flag in [True, False]:
    print("* Test with static train flag", train_flag)
    with make_scope() as session:
      network = TFNetwork(config=config, train_flag=train_flag)
      network.construct_from_dict(net_dict)
      network.initialize_params(session)
      feed_dict = make_feed_dict(network.extern_data)
      session.run(network.get_default_output_layer().output.placeholder, feed_dict=feed_dict)


def test_CondLayer_outside_layer_access():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim, ImplicitDynSizeDim, ImplicitSparseDim

  time_dim = SpatialDim('time')
  input_dim = FeatureDim('input', 10)

  config = Config(dict(extern_data={
    'data': {
      'dim_tags': (
        batch_dim,
        time_dim,
        input_dim
      ),
      'dtype': 'float32',
      'available_for_inference': True
    }
  }))

  labels_dim = FeatureDim('labels', 5)
  labels_1_dim = labels_dim + 1
  keys_dim = FeatureDim('keys', 2)
  qkv_dim_ = (2 * keys_dim) + input_dim
  num_heads_dim = SpatialDim('num_heads', 1)
  keys_h_dim = keys_dim.div_left(num_heads_dim)
  value_h_dim = input_dim.div_left(num_heads_dim)
  qkv_dim = 2 * keys_h_dim + value_h_dim
  time_kv_dim = SpatialDim('time:kv')
  time_kv_0_dim = SpatialDim('time:kv')

  net_dict = {
    'output': {
      'class': 'copy',
      'from': 'loop/output',
      'out_shape': {batch_dim, time_dim, ImplicitSparseDim(labels_1_dim)}
    },
    'encoder': {
      'class': 'subnetwork',
      'from': [],
      'subnetwork': {
        'self_att': {
          'class': 'subnetwork',
          'from': [],
          'subnetwork': {
            'qkv': {
              'class': 'subnetwork',
              'from': [],
              'subnetwork': {
                'weight': {
                  'class': 'variable',
                  'shape': [input_dim, qkv_dim_],
                  'param_name': 'param',
                },
                'bias': {
                  'class': 'variable',
                  'shape': [qkv_dim_],
                  'param_name': 'param',
                  'init': 0.0
                },
                'output': {
                  'class': 'copy',
                  'from': 'bias',
                  'out_shape': {qkv_dim_}
                }
              }
            },
            'output': {
              'class': 'copy',
              'from': 'qkv',
            }
          }
        },
        'output': {
          'class': 'copy',
          'from': 'self_att',
        }
      }
    },
    'out_label_logits': {
      'class': 'subnetwork',
      'from': [],
      'subnetwork': {
        'weight': {
          'class': 'variable',
          'shape': [input_dim, labels_1_dim],
          'param_name': 'param',
        },
        'bias': {
          'class': 'variable',
          'shape': [labels_1_dim],
          'param_name': 'param',
          'init': 0.0
        },
        'output': {
          'class': 'copy',
          'from': 'bias',
          'out_shape': {labels_1_dim}
        }
      }
    },
    'train_flag': {
      'class': 'train_flag',
      'out_shape': {}
    },
    'cond': {
      'class': 'cond',
      'from': [],
      'condition': 'train_flag',
      'true_layer': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'encoder': {
            'class': 'subnetwork',
            'from': [],
            'subnetwork': {
              'self_att': {
                'class': 'subnetwork',
                'from': [],
                'subnetwork': {
                  'qkv': {
                    'class': 'subnetwork',
                    'from': [],
                    'subnetwork': {
                      'dot': {
                        'class': 'dot',
                        'from': ['base:base:base:base:data:data', 'base:base:base:base:encoder/self_att/qkv/weight'],
                        'reduce': input_dim,
                        'out_shape': {batch_dim, time_dim, qkv_dim_}
                      },
                      'add': {
                        'class': 'combine',
                        'from': ['dot', 'base:base:base:base:encoder/self_att/qkv'],
                        'kind': 'add',
                        'out_shape': {batch_dim, time_dim, qkv_dim_}
                      },
                      'output': {
                        'class': 'copy',
                        'from': 'add',
                        'out_shape': {batch_dim, time_dim, qkv_dim_}
                      }
                    }
                  },
                  'qkv_split_dims': {
                    'class': 'split_dims',
                    'from': 'qkv',
                    'axis': qkv_dim_,
                    'dims': (num_heads_dim, qkv_dim),
                    'out_shape': {batch_dim, time_dim, num_heads_dim, qkv_dim}
                  },
                  'split': {
                    'class': 'split',
                    'from': 'qkv_split_dims',
                    'axis': qkv_dim,
                    'out_dims': (keys_h_dim, keys_h_dim, value_h_dim),
                    'out_shape': {batch_dim, time_dim, num_heads_dim, qkv_dim}
                  },
                  'reinterpret_new_dim': {
                    'class': 'reinterpret_data',
                    'set_dim_tags': {time_dim: time_kv_dim},
                    'from': 'split/1',
                    'out_shape': {batch_dim, num_heads_dim, keys_h_dim, time_kv_dim}
                  },
                  'reinterpret_new_dim_0': {
                    'class': 'reinterpret_data',
                    'set_dim_tags': {time_dim: time_kv_dim},
                    'from': 'split/2',
                    'out_shape': {batch_dim, num_heads_dim, value_h_dim, time_kv_dim}
                  },
                  'add': {
                    'class': 'copy',
                    'from': 'split/0',
                  },
                  'add_0': {
                    'class': 'copy',
                    'from': 'split/0',
                  },
                  'dot': {
                    'class': 'dot',
                    'from': ['add', 'reinterpret_new_dim'],
                    'reduce': keys_h_dim,
                    'out_shape': {batch_dim, time_dim, num_heads_dim, time_kv_dim}
                  },
                  'add_1': {
                    'class': 'copy',
                    'from': 'dot',
                  },
                  'mul': {
                    'class': 'eval',
                    'from': 'add_1',
                    'eval': 'source(0) * 0.7071067811865476',
                    'out_shape': {batch_dim, time_dim, num_heads_dim, time_kv_dim}
                  },
                  'att_weights': {
                    'class': 'softmax_over_spatial',
                    'from': 'mul',
                    'axis': time_kv_dim,
                    'out_shape': {batch_dim, time_dim, num_heads_dim, time_kv_dim}
                  },
                  'dropout': {
                    'class': 'dropout',
                    'from': 'att_weights',
                    'dropout': 0.1,
                    'dropout_axis': time_kv_dim,
                    'out_shape': {batch_dim, time_dim, num_heads_dim, time_kv_dim}
                  },
                  'att': {
                    'class': 'dot',
                    'from': ['dropout', 'reinterpret_new_dim_0'],
                    'reduce': time_kv_dim,
                    'out_shape': {batch_dim, time_dim, num_heads_dim, value_h_dim}
                  },
                  'output_0': {
                    'class': 'merge_dims',
                    'from': 'att',
                    'axes': (num_heads_dim, value_h_dim),
                    'out_dim': input_dim,
                    'out_shape': {batch_dim, time_dim, input_dim}
                  },
                  'output': {
                    'class': 'copy',
                    'from': 'output_0',
                    'out_shape': {batch_dim, time_dim, input_dim}
                  }
                }
              },
              'output': {
                'class': 'copy',
                'from': 'self_att',
                'out_shape': {batch_dim, time_dim, input_dim}
              }
            }
          },
          'output': {
            'class': 'copy',
            'from': 'encoder',
            'out_shape': {batch_dim, time_dim, input_dim}
          }
        }
      },
      'false_layer': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'encoder': {
            'class': 'subnetwork',
            'from': [],
            'subnetwork': {
              'self_att': {
                'class': 'subnetwork',
                'from': [],
                'subnetwork': {
                  'qkv': {
                    'class': 'subnetwork',
                    'from': [],
                    'subnetwork': {
                      'dot': {
                        'class': 'dot',
                        'from': ['base:base:base:base:data:data', 'base:base:base:base:encoder/self_att/qkv/weight'],
                        'reduce': input_dim,
                        'out_shape': {batch_dim, time_dim, qkv_dim_}
                      },
                      'add': {
                        'class': 'combine',
                        'from': ['dot', 'base:base:base:base:encoder/self_att/qkv'],
                        'kind': 'add',
                        'out_shape': {batch_dim, time_dim, qkv_dim_}
                      },
                      'output': {
                        'class': 'copy',
                        'from': 'add',
                        'out_shape': {batch_dim, time_dim, qkv_dim_}
                      }
                    }
                  },
                  'qkv_split_dims': {
                    'class': 'split_dims',
                    'from': 'qkv',
                    'axis': qkv_dim_,
                    'dims': (num_heads_dim, qkv_dim),
                    'out_shape': {batch_dim, time_dim, num_heads_dim, qkv_dim}
                  },
                  'split': {
                    'class': 'split',
                    'from': 'qkv_split_dims',
                    'axis': qkv_dim,
                    'out_dims': (keys_h_dim, keys_h_dim, value_h_dim),
                    'out_shape': {batch_dim, time_dim, num_heads_dim, qkv_dim}
                  },
                  'reinterpret_new_dim': {
                    'class': 'reinterpret_data',
                    'set_dim_tags': {time_dim: time_kv_0_dim},
                    'from': 'split/1',
                    'out_shape': {batch_dim, num_heads_dim, keys_h_dim, time_kv_0_dim}
                  },
                  'reinterpret_new_dim_0': {
                    'class': 'reinterpret_data',
                    'set_dim_tags': {time_dim: time_kv_0_dim},
                    'from': 'split/2',
                    'out_shape': {batch_dim, num_heads_dim, value_h_dim, time_kv_0_dim}
                  },
                  'add': {
                    'class': 'copy',
                    'from': 'split/0',
                  },
                  'add_0': {
                    'class': 'copy',
                    'from': 'split/0',
                  },
                  'dot': {
                    'class': 'dot',
                    'from': ['add', 'reinterpret_new_dim'],
                    'reduce': keys_h_dim,
                    'out_shape': {batch_dim, time_dim, num_heads_dim, time_kv_0_dim}
                  },
                  'add_1': {
                    'class': 'copy',
                    'from': 'dot',
                  },
                  'mul': {
                    'class': 'eval',
                    'from': 'add_1',
                    'eval': 'source(0) * 0.7071067811865476',
                    'out_shape': {batch_dim, time_dim, num_heads_dim, time_kv_0_dim}
                  },
                  'att_weights': {
                    'class': 'softmax_over_spatial',
                    'from': 'mul',
                    'axis': time_kv_0_dim,
                    'out_shape': {batch_dim, time_dim, num_heads_dim, time_kv_0_dim}
                  },
                  'dropout': {
                    'class': 'dropout',
                    'from': 'att_weights',
                    'dropout': 0.1,
                    'dropout_axis': time_kv_0_dim,
                    'out_shape': {batch_dim, time_dim, num_heads_dim, time_kv_0_dim}
                  },
                  'att': {
                    'class': 'dot',
                    'from': ['dropout', 'reinterpret_new_dim_0'],
                    'reduce': time_kv_0_dim,
                    'out_shape': {batch_dim, time_dim, num_heads_dim, value_h_dim}
                  },
                  'output_0': {
                    'class': 'merge_dims',
                    'from': 'att',
                    'axes': (num_heads_dim, value_h_dim),
                    'out_dim': input_dim,
                    'out_shape': {batch_dim, time_dim, input_dim}
                  },
                  'output': {
                    'class': 'copy',
                    'from': 'output_0',
                    'out_shape': {batch_dim, time_dim, input_dim}
                  }
                }
              },
              'output': {
                'class': 'copy',
                'from': 'self_att',
                'out_shape': {batch_dim, time_dim, input_dim}
              }
            }
          },
          'output': {
            'class': 'copy',
            'from': 'encoder',
            'out_shape': {batch_dim, time_dim, input_dim}
          }
        }
      },
      'out_shape': {batch_dim, time_dim, input_dim},
      'name_scope': ''
    },
    'loop': {
      'class': 'rec',
      'from': [],
      'unit': {
        'rec_unstack': {
          'class': 'rec_unstack',
          'from': 'base:cond',
          'axis': time_dim,
          'out_shape': {batch_dim, input_dim}
        },
        'model': {
          'class': 'subnetwork',
          'from': [],
          'subnetwork': {
            'out_label_logits': {
              'class': 'subnetwork',
              'from': [],
              'subnetwork': {
                'dot': {
                  'class': 'dot',
                  'from': ['base:base:rec_unstack', 'base:base:base:out_label_logits/weight'],
                  'reduce': input_dim,
                  'out_shape': {batch_dim, labels_1_dim}
                },
                'add': {
                  'class': 'combine',
                  'from': ['dot', 'base:base:base:out_label_logits'],
                  'kind': 'add',
                  'out_shape': {batch_dim, labels_1_dim}
                },
                'output': {
                  'class': 'copy',
                  'from': 'add',
                  'out_shape': {batch_dim, labels_1_dim}
                }
              }
            },
            'output': {
              'class': 'copy',
              'from': 'out_label_logits',
              'out_shape': {batch_dim, labels_1_dim}
            }
          },
          'name_scope': ''
        },
        'softmax': {
          'class': 'softmax_over_spatial',
          'from': 'model',
          'axis': labels_1_dim,
          'log_space': True,
          'out_shape': {batch_dim, labels_1_dim}
        },
        'choice': {
          'class': 'choice',
          'from': 'softmax',
          'target': None,
          'beam_size': 3,
          'search': True,
          'input_type': 'log_prob',
          'length_normalization': False,
          'initial_output': 'base:constant',
          'out_shape': {batch_dim, ImplicitSparseDim(labels_1_dim)}
        },
        'output': {
          'class': 'copy',
          'from': 'choice',
          'out_shape': {batch_dim, ImplicitSparseDim(labels_1_dim)}
        }
      },
      'axis': time_dim,
      'out_shape': {batch_dim, time_dim, ImplicitSparseDim(labels_1_dim)},
      'name_scope': ''
    },
    'length': {
      'class': 'length',
      'from': ['data:data'],
      'axis': time_dim,
      'out_shape': {batch_dim}
    },
    'mul': {
      'class': 'eval',
      'from': 'encoder',
      'eval': 'source(0) * 2',
      'out_shape': {}
    },
    'constant': {
      'class': 'constant',
      'value': 5,
      'shape': [batch_dim],
      'sparse_dim': labels_1_dim,
      'shape_deps': ['data:data']
    }
  }
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    net.initialize_params(session=session)
    out = net.get_default_output_layer()
    session.run(out.output.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_CondLayer_multiple_outputs():
  net_dict = {
    "cond": {
      "class": "cond", "from": [],
      "condition": "data:cond",
      "true_layer": {
        "class": "subnetwork", "from": [], "subnetwork": {
          "a": {"class": "constant", "value": 2, "is_output_layer": True},
          "b": {"class": "constant", "value": 3, "is_output_layer": True},
          "output": {"class": "copy", "from": "a"}
        }
      },
      "false_layer": {
        "class": "subnetwork", "from": [], "subnetwork": {
          "a": {"class": "constant", "value": 5, "is_output_layer": True},
          "b": {"class": "constant", "value": 7, "is_output_layer": True},
          "output": {"class": "copy", "from": "a"}
        }
      }
    },
    "output": {
      "class": "combine", "kind": "mul", "from": ["cond/a", "cond/b"]
    }
  }
  config = Config(dict(extern_data={"cond": {"dim_tags": (), "dtype": "bool", "available_for_inference": True}}))
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    out_t = net.get_default_output_layer().output.placeholder
    assert_equal(session.run(out_t, feed_dict={net.extern_data.data["cond"].placeholder: True}), 6)
    assert_equal(session.run(out_t, feed_dict={net.extern_data.data["cond"].placeholder: False}), 35)


def test_ScatterNdLayer_RangeLayer():
  from returnn.tf.util.data import batch_dim, Dim
  n_batch, n_time, n_ts, n_out = 2, 3, 6, 11
  time_dim = SpatialDim("T")
  feat_dim = FeatureDim("F", dimension=7)
  ts_dim = SpatialDim("ts", dimension=n_ts)
  rnd = numpy.random.RandomState(42)
  config = Config({
    "debug_print_layer_output_template": True,
    "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}}
  })
  net_dict = {
    "t": {"class": "eval", "from": [], "eval": "tf.convert_to_tensor([1, 2])",
          "out_type": {"shape": (), "dtype": "int32"}},  # (B,)
    "range": {"class": "range", "limit": n_ts, "out_spatial_dim": ts_dim},  # (Ts,)
    "add_t": {"class": "combine", "kind": "add", "from": ["t", "range"], "out_shape": {batch_dim, ts_dim}},  # (B,Ts)
    "t_rel_var": {"class": "variable", "shape": (ts_dim, n_out), "init": "glorot_uniform"},  # (B,Ts,D)
    "output": {"class": "scatter_nd", "from": "t_rel_var", "position": "add_t", "position_axis": ts_dim,
               "output_dim_via_time_from": "data", "filter_invalid_indices": True}
  }
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)

    fetches = network.get_fetches_dict()
    data_input = network.extern_data.data["data"]
    out_layer = network.get_default_output_layer()
    assert isinstance(out_layer, ScatterNdLayer)
    assert out_layer.output.shape == (None, 11)
    assert out_layer.output.feature_dim_axis_or_unspecified is NotSpecified and out_layer.output.feature_dim_axis == 2
    assert out_layer.output.time_dim_axis == 1

    session.run(tf_compat.v1.variables_initializer(tf_compat.v1.global_variables() + [network.global_train_step_var]))
    info, out = session.run(
      (fetches, out_layer.output.placeholder),
      feed_dict={
        data_input.batch.dim: n_batch,
        data_input.placeholder: rnd.normal(size=(n_batch, n_time, feat_dim.dimension)).astype("float32"),
        data_input.size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
      })
    print(info)
    print(out)  # random...


def _run_repeat_layer(session, net, input_data_layer):
  repetitions_data_layer = InternalLayer(
    name="repetitions", network=net,
    output=Data(**{'name': 'repetitions', 'shape': (None,), 'feature_dim_axis': None, 'dtype': 'int32'}))
  repetitions_data_layer.output.placeholder = tf.constant(
    [[1, 3, 2, 1, 3, 4, 1, 1, 2, 1],
     [3, 2, 1, 3, 0, 1, 1, 0, 0, 1]], dtype="int32")  # [B, T] (sparse), additional 1 in second sequence to test masking
  repetitions_data_layer.output.size_placeholder = {0: tf.constant([10, 7], dtype="int32")}  # [B]

  opts = {
    'network': net, 'name': 'repeat_layer_test', 'sources': [input_data_layer],
    'repetitions': repetitions_data_layer, 'axis': 'T'}
  out_data = RepeatLayer.get_out_data_from_opts(**opts)
  out_data.sanity_check()
  print(out_data)
  repeat_layer = RepeatLayer(output=out_data, **opts)
  print(repeat_layer.output)

  output, size_placeholder = session.run(
    [repeat_layer.output.placeholder, repeat_layer.output.size_placeholder.as_dict()],
    feed_dict=make_feed_dict(net.extern_data, n_batch=2))
  assert numpy.all(numpy.equal(size_placeholder[0], numpy.asarray([19, 11])))
  assert numpy.all(numpy.equal(output.shape, numpy.asarray([2, 19, 5])))
  # the 6 last positions of the second sequence need to be padded with zeros
  assert numpy.all(numpy.equal(output[1, 11:], 0))
  assert out_data.shape == (None, 5)
  assert out_data.batch_dim_axis == 0
  assert out_data.time_dim_axis == 1


def test_RepeatLayerBTF():
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData({"data": {"shape": (None, 3)}}))
    input_data_layer = InternalLayer(name="src", network=net, output=Data(name="src", shape=(None, 5), dim=5))
    input_data_layer.output.size_placeholder = {0: tf.constant([10, 7])}  # [B]
    input_data_layer.output.placeholder = tf_compat.v1.random_uniform((2, 10, 5))  # [B, T, F]

    _run_repeat_layer(session, net, input_data_layer)


def test_RepeatLayerTBF():
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData({"data": {"shape": (None, 3)}}))
    input_data_layer = InternalLayer(
      name="src", network=net,
      output=Data(**{'name': "src", 'shape': (None, 5), 'dim': 5, 'batch_dim_axis': 1, 'time_dim_axis': 0}))
    input_data_layer.output.size_placeholder = {0: tf.constant([10, 7])}  # [B]
    input_data_layer.output.placeholder = tf_compat.v1.random_uniform((10, 2, 5))  # [T, B, F]

    _run_repeat_layer(session, net, input_data_layer)


def test_RepeatLayerBFT():
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData({"data": {"shape": (None, 3)}}))
    input_data_layer = InternalLayer(
      name="src", network=net,
      output=Data(**{'name': 'src', 'shape': (5, None), 'dim': 5, 'time_dim_axis': 2, 'feature_dim_axis': 1}))
    input_data_layer.output.size_placeholder = {1: tf.constant([10, 7])}  # [B]
    input_data_layer.output.placeholder = tf_compat.v1.random_uniform((2, 5, 10))  # [B, F, T]

    _run_repeat_layer(session, net, input_data_layer)


def test_RepeatLayerTF():
  with make_scope() as session:
    # need to provide extern_data here to provide batch info
    net = TFNetwork(extern_data=ExternData(data={"data": {"dim": 2, "sparse": True}}), train_flag=True)
    input_data_layer = InternalLayer(
      name="src", network=net,
      output=Data(**{'name': 'src', 'shape': (None, 5), 'dim': 5, 'batch_dim_axis': None, 'time_dim_axis': 0}))
    input_data_layer.output.size_placeholder = {0: tf.constant([10])}  # []
    input_data_layer.output.placeholder = tf_compat.v1.random_uniform((10, 5))  # [T, F]

    _run_repeat_layer(session, net, input_data_layer)


def test_RepeatLayer_int_repetitions():
  with make_scope() as session:
    n_out = 5
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "output": {"class": "repeat", "repetitions": 3, "axis": "F", "from": ["data"]}
    })
    session.run(tf_compat.v1.global_variables_initializer())
    out = net.layers["output"].output
    n_batch = 3
    max_seq_len = 10
    feed = make_feed_dict(net.extern_data.data.values(), n_batch=n_batch, n_time=max_seq_len, same_time=True)
    v = session.run(out.placeholder, feed_dict=feed)
    input_data = feed[net.extern_data.data["data"].placeholder]

    assert out.batch_dim_axis == 0
    ref = numpy.swapaxes(numpy.repeat(input_data, 3, axis=-1), -1, out.feature_dim_axis)
    numpy.testing.assert_allclose(ref, v, rtol=1e-5)


def test_RepeatLayer_int():
  # https://github.com/rwth-i6/returnn_common/issues/162
  from returnn.tf.util.data import batch_dim, SpatialDim
  time_dim = SpatialDim('time')
  config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim), "dtype": "float32"}}})
  _repeat_out_dim = time_dim * 5
  net_dict = {
    'output': {
      'class': 'repeat',
      'from': 'data:data',
      'repetitions': 5,
      'axis': time_dim,
      'out_dim': _repeat_out_dim,
      'out_shape': {batch_dim, _repeat_out_dim}
    },
  }
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    in_ = net.extern_data.data["data"]
    out = net.get_default_output_layer().output
    _, in_seq_len, _, out_seq_len = session.run(
      (in_.placeholder, in_.get_sequence_lengths(), out.placeholder, out.get_sequence_lengths()),
      feed_dict=make_feed_dict(net.extern_data))
    assert (in_seq_len * 5 == out_seq_len).all()


def test_RepeatLayer_as_loss_flatten():
  # https://github.com/rwth-i6/returnn_common/issues/201
  # https://github.com/rwth-i6/returnn/issues/1115
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim('time')
  in_dim = FeatureDim('in', 3)
  config = Config(dict(extern_data={'data': {'dim_tags': (batch_dim, time_dim, in_dim)}}))
  net_dict = {
    'output': {
      'class': 'repeat',
      'from': 'data:data',
      'repetitions': 1,
      'axis': time_dim,
      'out_dim': time_dim,
      'loss': 'as_is',
      'out_shape': {batch_dim, time_dim, in_dim}
    },
  }
  with make_scope() as session:
    net = TFNetwork(config=config, eval_flag=True)
    net.construct_from_dict(net_dict)
    feed_dict = make_feed_dict(net.extern_data)
    fetches = net.get_fetches_dict(should_eval=True)
    session.run(fetches, feed_dict=feed_dict)


def test_TileLayer():
  with make_scope() as session:
    n_out = 5
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "output": {"class": "tile", "multiples": {"F": 3}, "from": ["data"]}
    })
    session.run(tf_compat.v1.global_variables_initializer())
    out = net.layers["output"].output.placeholder
    n_batch = 3
    max_seq_len = 10
    feed = make_feed_dict(net.extern_data.data.values(), n_batch=n_batch, n_time=max_seq_len, same_time=True)
    v = session.run(out, feed_dict=feed)
    input_len = feed[net.extern_data.data["data"].size_placeholder[0]]
    input_data = feed[net.extern_data.data["data"].placeholder]

    ref = numpy.tile(input_data, [1, 1, 3])

    numpy.testing.assert_allclose(ref, v, rtol=1e-5)


def test_ScatterNdLayer_RangeLayer_RangeInAxisLayer():
  from returnn.tf.util.data import batch_dim, Dim, ImplicitDynSizeDim
  n_batch, n_time, n_ts, n_in, n_out = 2, 3, 6, 7, 11
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("in-feature", dimension=n_in)
  ts_dim = SpatialDim("ts", dimension=n_ts)
  rnd = numpy.random.RandomState(42)
  config = Config({
    "debug_print_layer_output_template": True,
    "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}}
  })
  net_dict = {
    "t": {
      "class": "range_in_axis", "axis": "t", "from": "data",
      "out_shape": {time_dim, ImplicitDynSizeDim(batch_dim)}},  # (T,)
    "range": {"class": "range", "limit": n_ts, "out_spatial_dim": ts_dim},  # (Ts,)
    "add_t": {
      "class": "combine", "kind": "add", "from": ["t", "range"],
      "out_shape": {time_dim, ts_dim, ImplicitDynSizeDim(batch_dim)}},  # (T,Ts)
    "t_rel_var": {"class": "variable", "shape": (ts_dim, n_out), "init": "glorot_uniform"},  # (Ts,D)
    "output": {"class": "scatter_nd", "from": "t_rel_var", "position": "add_t", "position_axis": ts_dim,
               "output_dim_via_time_from": "data", "filter_invalid_indices": True}  # (T,T,D)
  }
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)

    t_layer = network.layers["t"]
    assert isinstance(t_layer, RangeInAxisLayer)
    assert t_layer.output.time_dim_axis == 0

    fetches = network.get_fetches_dict()
    data_input = network.extern_data.data["data"]
    out_layer = network.get_default_output_layer()
    assert isinstance(out_layer, ScatterNdLayer)
    assert out_layer.output.shape == (None, None, 11)
    assert out_layer.output.feature_dim_axis_or_unspecified is NotSpecified and out_layer.output.feature_dim_axis == 2
    assert out_layer.output.time_dim_axis == 0

    session.run(tf_compat.v1.variables_initializer(tf_compat.v1.global_variables() + [network.global_train_step_var]))
    info, out = session.run(
      (fetches, out_layer.output.placeholder),
      feed_dict={
        data_input.placeholder: rnd.normal(size=(n_batch, n_time, n_in)).astype("float32"),
        data_input.size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
      })
    print(info)
    print(out)  # random...


def test_ScatterNdLayer_pos_batch_last_dim():
  config = Config({
    "debug_print_layer_output_template": True,
    "extern_data": {"data": {"dim": 13}}
  })
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    data = network.construct_layer({}, "data")
    pos = InternalLayer(
      name="pos", network=network,
      output=Data(
        name='pos', shape=(None, 6), dtype='int32', sparse=True, dim=None, batch_dim_axis=2,
        auto_create_placeholders=True))
    val = InternalLayer(
      name="val", network=network,
      output=Data(
        name='var', shape=(6, 11), time_dim_axis=None,
        auto_create_placeholders=True))
    scatter_opts = dict(
      name="scatter", network=network,
      sources=[val], position=pos, position_axis="dim:6",
      output_dim_via_time_from=data, filter_invalid_indices=True)
    scatter_out_template = ScatterNdLayer.get_out_data_from_opts(**scatter_opts)
    print("scatter out:", scatter_out_template)
    assert scatter_out_template.shape == (None, None, 11) and scatter_out_template.batch_ndim == 4
    scatter = ScatterNdLayer(output=scatter_out_template, **scatter_opts)
    print("scatter out dim tags:")
    pprint(scatter.output.get_batch_shape_dim_tags())
    assert_equal(scatter.output.get_size_dim_tag(0), pos.output.get_time_dim_tag())
    assert_equal(scatter.output.get_size_dim_tag(1), data.output.get_time_dim_tag())
    session.run(scatter.output.placeholder, feed_dict=make_feed_dict([data.output, pos.output, val.output]))


def test_ConvLayer_get_valid_out_dim():
  assert_equal(ConvLayer.calc_out_dim(in_dim=10, stride=1, filter_size=2, padding="same"), 10)
  assert_equal(ConvLayer.calc_out_dim(in_dim=10, stride=1, filter_size=3, padding="same"), 10)
  assert_equal(ConvLayer.calc_out_dim(in_dim=10, stride=1, filter_size=2, padding="valid"), 9)
  assert_equal(ConvLayer.calc_out_dim(in_dim=10, stride=1, filter_size=3, padding="valid"), 8)
  assert_equal(ConvLayer.calc_out_dim(in_dim=10, stride=2, filter_size=2, padding="valid"), 5)
  assert_equal(ConvLayer.calc_out_dim(in_dim=10, stride=3, filter_size=2, padding="valid"), 3)
  assert_equal(ConvLayer.calc_out_dim(in_dim=10, stride=3, filter_size=1, padding="valid"), 4)
  assert_equal(ConvLayer.calc_out_dim(in_dim=10, stride=3, filter_size=2, padding="same"), 4)
  assert_equal(ConvLayer.calc_out_dim(in_dim=41, stride=1, filter_size=2, padding="valid"), 40)
  assert_equal(ConvLayer.calc_out_dim(in_dim=40, stride=2, filter_size=2, padding="valid"), 20)
  assert_equal(ConvLayer.calc_out_dim(in_dim=2, stride=1, filter_size=3, padding="valid"), 0)


def test_LengthLayer():
  net_dict = {
    "output": {"class": "length", "from": "data", "axis": "T"},
  }
  with make_scope() as session:
    config = Config({"extern_data": {"data": {"dim": 10}}})
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    in_ = net.extern_data.get_default_input_data()
    out = net.get_default_output_layer().output
    in_v, in_size_v, out_v = session.run(
      (in_.placeholder, in_.size_placeholder[0], out.placeholder), feed_dict=make_feed_dict(net.extern_data))
    n_batch, n_time, n_feat = in_v.shape
    assert out_v.shape == in_size_v.shape == (n_batch,)
    assert list(out_v) == list(in_size_v)


def test_LengthLayer_batch():
  net_dict = {
    "input_flat": {"class": "flatten_batch", "from": "data"},  # [B_T,F]
    "output": {"class": "length", "from": "input_flat", "axis": "B"},  # scalar -> B_T
  }
  with make_scope() as session:
    config = Config({"extern_data": {"data": {"dim": 10}}})
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    in_ = net.extern_data.get_default_input_data()
    flat = net.layers["input_flat"].output
    out = net.get_default_output_layer().output
    in_shape_v, in_lens_v, flat_shape_v, out_v = session.run(
      (tf.shape(in_.placeholder), in_.get_sequence_lengths(),
       tf.shape(flat.placeholder), out.placeholder),
      feed_dict=make_feed_dict(net.extern_data))
    assert in_lens_v.shape == (in_shape_v[0],)
    assert max(in_lens_v) == in_shape_v[1]
    assert sum(in_lens_v) == flat_shape_v[0] == out_v


def test_LengthLayer_generic_dim():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  out_time_dim = time_dim + 2
  feat_dim = FeatureDim("feat", 3)
  config = Config({"extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}}})
  net_dict = {
    "output": {"class": "length", "from": "data", "axis": out_time_dim},
  }
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    in_ = net.extern_data.get_default_input_data()
    out = net.get_default_output_layer().output
    in_v, in_size_v, out_v = session.run(
      (in_.placeholder, in_.size_placeholder[0], out.placeholder), feed_dict=make_feed_dict(net.extern_data))
    n_batch, n_time, n_feat = in_v.shape
    assert out_v.shape == in_size_v.shape == (n_batch,)
    assert list(out_v) == list(in_size_v + 2)


def test_RandIntLayer():
  with make_scope() as session:
    from returnn.tf.util.data import Dim
    n_out = 5
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out}}
    })
    net = TFNetwork(config=config, train_flag=True)
    n_batch = 3
    max_seq_len = 10
    feed = make_feed_dict(net.extern_data.data.values(), n_batch=n_batch, n_time=max_seq_len, same_time=True)
    size_placeholder = net.extern_data.data["data"].size_placeholder[0]
    input_len = feed[size_placeholder]
    sz = (
      Dim(description="feature", kind=Dim.Types.Feature, dimension=5),
      Dim(kind=Dim.Types.Batch),
      net.extern_data.data["data"].get_size_dim_tag(0),
      3,
    )
    net.construct_from_dict({
      "output": {"class": "rand_int", "shape": sz, "minval": 3, "maxval": 10, "seed": 42}
    })
    session.run(tf_compat.v1.global_variables_initializer())
    out = net.layers["output"].output.placeholder
    v = session.run(out, feed_dict=feed)

    assert_equal(v.shape, (5, n_batch, max(input_len), 3))


def test_rand_indices():
  with make_scope() as session:
    from returnn.tf.util.data import FeatureDim, SpatialDim, batch_dim
    feature_dim = FeatureDim("feature", 5)
    time_dim = SpatialDim("time")
    config = Config({
      "extern_data": {"data": {"dim_tags": (batch_dim, time_dim, feature_dim)}}
    })
    net = TFNetwork(config=config, train_flag=True)
    sz = (batch_dim, time_dim, SpatialDim("other-spatial", 7))
    net.construct_from_dict({
      "flat": {"class": "flatten_batch", "from": "data"},
      "length_flat": {"class": "length", "from": "flat", "axis": batch_dim},
      "indices_flat": {"class": "rand_int", "shape": sz, "minval": 0, "maxval": "length_flat", "seed": 42},
      "output": {"class": "gather", "from": "flat", "axis": batch_dim, "position": "indices_flat"},
    })

    n_batch, n_time = 3, 11
    indices_flat, output = session.run(
      (net.layers["indices_flat"].output.placeholder, net.layers["output"].output.placeholder),
      feed_dict=make_feed_dict(net.extern_data, n_batch=n_batch, n_time=n_time))
    assert_equal(indices_flat.shape, (n_batch, n_time, sz[-1].dimension))
    assert_equal(output.shape, (n_batch, n_time, sz[-1].dimension, feature_dim.dimension))


def test_RandomLayer():
  # https://github.com/rwth-i6/returnn_common/issues/197
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim('time')
  input_dim = FeatureDim('input', 3)
  config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, input_dim)}}})
  net_dict = {
    'random': {
      'class': 'random',
      'shape': [batch_dim, input_dim],
      'distribution': 'normal',
      'mean': 0.0,
      'stddev': 1.0,
    },
    'output': {
      'class': 'combine', "kind": "add",
      'from': ["data:data", 'random'],
      'out_shape': {batch_dim, time_dim, input_dim}
    }
  }
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    net.initialize_params(session)
    out = net.get_default_output_layer().output.copy_as_time_major()
    out_np = session.run(out.placeholder, feed_dict=make_feed_dict(net.extern_data))
    print(out_np)


def test_RandomLayer_shape_deps():
  # https://github.com/rwth-i6/returnn_common/issues/197
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim('time')
  time_pool_dim = SpatialDim("time-pool")
  input_dim = FeatureDim('input', 3)
  config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, input_dim)}}})
  net_dict = {
    'pool': {"class": "pool", "pool_size": [2], "mode": "max", "from": "data", "out_spatial_dims": [time_pool_dim]},
    'random': {
      'class': 'random',
      'shape': [batch_dim, time_pool_dim, input_dim],
      "shape_deps": ["pool"],
      'distribution': 'normal',
      'mean': 0.0,
      'stddev': 1.0,
    },
    'output': {
      'class': 'combine', "kind": "add",
      'from': ['random', 'pool'],  # the order is relevant to test shape_deps
      'out_shape': {batch_dim, time_pool_dim, input_dim}
    }
  }
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    net.initialize_params(session)
    out = net.get_default_output_layer().output.copy_as_time_major()
    out_np = session.run(out.placeholder, feed_dict=make_feed_dict(net.extern_data))
    print(out_np)


def test_RandomLayer_in_loop():
  # https://github.com/rwth-i6/returnn/issues/1044
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim('time')
  input_dim = FeatureDim('input', 3)
  config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, input_dim)}}})
  random_state_dim = FeatureDim('random-state', 3)
  net_dict = {
    'output': {
      'class': 'rec',
      'from': "data",
      'unit': {
        'rnd': {
          'class': 'subnetwork',
          'from': [],
          'subnetwork': {
            'random': {
              'class': 'random',
              'shape': [
                batch_dim,
                input_dim
              ],
              'distribution': 'normal',
              'mean': 0.0,
              'stddev': 1.0,
              'explicit_state': 'base:base:rnd_state_var0',
              'auto_update_state': True
            },
            'output': {
              'class': 'copy',
              'from': 'random',
              'out_shape': {batch_dim, input_dim}
            }
          },
          'out_shape': {batch_dim, input_dim}
        },
        'output': {
          'class': 'eval',
          'from': ['data:source', 'rnd'],
          'eval': 'source(0) * 0.0 + source(1)',
          'out_shape': {batch_dim, input_dim}
        },
      },
      'axis': time_dim,
      'out_shape': {batch_dim, time_dim, input_dim},
      'name_scope': ''
    },
    'random_state_init': {
      'class': 'random_state_init',
      'out_dim': random_state_dim,
      'out_shape': {random_state_dim}
    },
    'rnd_state_var0': {
      'class': 'variable',
      'shape': [
        random_state_dim
      ],
      'param_name': 'param',
      'dtype': 'int64',
      'init_by_layer': 'random_state_init',
      'name_scope': 'rnd/state_var0'
    }
  }
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    net.initialize_params(session)
    out = net.get_default_output_layer().output.copy_as_time_major()
    out_np = session.run(out.placeholder, feed_dict=make_feed_dict(net.extern_data))
    print(out_np)
    out0_np = out_np[:1]  # [1,B,D]
    print((out0_np == out_np))
    assert not (out0_np == out_np).all()  # not all the same


def test_RandomLayer_zero_shape():
  # https://github.com/rwth-i6/returnn/issues/1190
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim('time')
  input_dim = FeatureDim('input', 3)
  config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, input_dim)}}})
  net_dict = {
    'random': {
      'class': 'random',
      'shape': [batch_dim, time_dim, FeatureDim("zero", 0)],
      "shape_deps": ["data"],
      'distribution': 'normal',
      'mean': 0.0,
      'stddev': 1.0,
    },
    'output': {
      'class': 'copy', 'from': ['data', 'random'],
    }
  }
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    net.initialize_params(session)
    in_ = net.extern_data.get_default_input_data()
    out = net.get_default_output_layer().output
    in_np, out_np = session.run((in_.placeholder, out.placeholder), feed_dict=make_feed_dict(net.extern_data))
    print(in_np.shape, out_np.shape)
    print(out_np)
    assert in_np.shape == out_np.shape


def test_untrainable_params():
  with make_scope() as session:
    config = Config()
    n_in, n_out = 2, 3
    config.update({
      "num_outputs": n_out,
      "num_inputs": n_in,
      "network": {
        "l1": {"class": "linear", "activation": None, "n_out": n_out, "from": "data:data"},
        "output": {"class": "linear", "activation": None, "from": ["l1"], "n_out": n_out, "trainable": False}
      }
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    l1 = network.layers["l1"]
    l2 = network.layers["output"]
    assert_equal(set(network.get_trainable_params()), {l1.params["W"], l1.params["b"]})


def test_reuse_params():
  with make_scope() as session:
    config = Config()
    n_in, n_out = 2, 3
    config.update({
      "num_outputs": n_out,
      "num_inputs": n_in,
      "network": {
        "l1": {"class": "linear", "activation": None, "n_out": n_out, "from": "data:data"},
        "output": {"class": "linear", "activation": None, "n_out": n_out, "from": "data:data", "reuse_params": "l1"}
      }
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    l1 = network.layers["l1"]
    l2 = network.layers["output"]
    assert_equal(set(l1.params.keys()), {"W", "b"})
    assert_equal(set(l2.params.keys()), set())
    assert_equal(set(network.get_trainable_params()), {l1.params["W"], l1.params["b"]})


def test_reuse_params_map_custom():
  with make_scope() as session:
    config = Config()
    n_in, n_out = 2, 3
    config.update({
      "num_outputs": n_out,
      "num_inputs": n_in,
      "network": {
        "l1": {"class": "linear", "activation": "tanh", "with_bias": False, "n_out": 5, "from": "data:data"},
        "output": {
          "class": "linear", "activation": None, "n_out": n_in, "from": ["l1"], "target": "data",
          "reuse_params": {
            "map": {
              "W": {
                "reuse_layer": "l1",
                "custom": (lambda reuse_layer, **kwargs: tf.transpose(reuse_layer.params["W"]))},
              "b": None}
          },
        }
      }
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    l1 = network.layers["l1"]
    l2 = network.layers["output"]
    assert_equal(set(l1.params.keys()), {"W"})
    assert_equal(set(l2.params.keys()), {"b"})
    assert_equal(set(network.get_trainable_params()), {l1.params["W"], l2.params["b"]})


def test_reuse_params_map_custom_rev():
  with make_scope() as session:
    config = Config()
    n_in, n_out = 2, 3
    config.update({
      "num_outputs": n_out,
      "num_inputs": n_in,
      "network": {
        "output": {"class": "linear", "activation": "tanh", "with_bias": False, "from": ["l1"], "n_out": n_in},
        "l1": {
          "class": "linear", "activation": None, "n_out": 5, "from": ["data"], "target": "data",
          "reuse_params": {
            "map": {
              "W": {
                "reuse_layer": "output",
                "custom": (lambda reuse_layer, **kwargs: tf.transpose(reuse_layer.params["W"]))},
              "b": None}
          },
        }
      }
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    l1 = network.layers["l1"]
    l2 = network.layers["output"]
    assert_equal(set(l1.params.keys()), {"b"})
    assert_equal(set(l2.params.keys()), {"W"})
    assert_equal(set(network.get_trainable_params()), {l2.params["W"], l1.params["b"]})


def test_reuse_params_map_custom_dep_loop():
  config = Config()
  n_in, n_out = 2, 3
  config.update({
    "num_outputs": n_out,
    "num_inputs": n_in,
    "network": {
      "encoder": {"class": "copy", "from": ["data"]},
      "enc_ctx": {"class": "linear", "activation": None, "with_bias": True, "from": ["encoder"], "n_out": 10},
      "inv_fertility": {"class": "linear", "activation": "sigmoid", "with_bias": False, "from": ["encoder"],
                        "n_out": 1},
      "output": {"class": "rec", "from": [], "unit": {
        'output': {'class': 'choice', 'target': 'classes', 'beam_size': 5, 'from': ["output_prob"],
                   "initial_output": 0},
        "end": {"class": "compare", "from": ["output"], "value": 0},
        'target_embed': {'class': 'linear', 'activation': None, "with_bias": False, 'from': ['output'], "n_out": 6,
                         "initial_output": 0},
        "weight_feedback": {"class": "linear", "activation": None, "with_bias": False,
                            "from": ["prev:accum_att_weights"], "n_out": 10},
        "prev_s_state": {"class": "get_last_hidden_state", "from": ["prev:s"], "n_out": 20},
        "prev_s_transformed": {"class": "linear", "activation": None, "with_bias": False, "from": ["prev_s_state"],
                               "n_out": 10},
        "energy_in": {"class": "combine", "kind": "add",
                      "from": ["base:enc_ctx", "weight_feedback", "prev_s_transformed"], "n_out": 10},
        "energy_tanh": {"class": "activation", "activation": "tanh", "from": ["energy_in"]},
        "energy": {"class": "linear", "activation": None, "with_bias": False, "from": ["energy_tanh"], "n_out": 1},
        "att_weights": {"class": "softmax_over_spatial", "from": ["energy"]},
        "accum_att_weights": {"class": "eval",
                              "from": ["prev:accum_att_weights", "att_weights", "base:inv_fertility"],
                              "eval": "source(0) + source(1) * source(2) * 0.5",
                              "out_type": {"dim": 1, "shape": (None, 1)}},
        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
        "s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["target_embed", "att"], "n_out": 10},
        "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
                       "n_out": 2 * 6},
        "readout": {"class": "reduce_out", "mode": "max", "num_pieces": 2, "from": ["readout_in"]},
        "output_prob": {
          "class": "softmax", "from": ["readout"], "dropout": 0.3,
          "reuse_params": {
            "map": {
              "W": {
                "reuse_layer": "target_embed",
                "custom": (lambda reuse_layer, **kwargs: tf.transpose(reuse_layer.params["W"]))},
              "b": None}},
          "target": "classes", "loss": "ce", "loss_opts": {"label_smoothing": 0.1}}
      }, "target": "classes", "max_seq_len": "max_len_from('base:encoder')"},
    }
  })
  with make_scope() as session:
    print("Construct for training")
    from returnn.tf.layers.rec import RecLayer, _SubnetworkRecCell
    train_net = TFNetwork(config=config, train_flag=True)
    train_net.construct_from_dict(config.typed_dict["network"])
    train_rec_layer = train_net.layers["output"]
    assert isinstance(train_rec_layer, RecLayer)
    assert isinstance(train_rec_layer.cell, _SubnetworkRecCell)
    assert_equal(set(train_rec_layer.cell.input_layers_moved_out), {"output", "target_embed"})
    assert_equal(set(train_rec_layer.cell.output_layers_moved_out), {"output_prob", "readout", "readout_in"})
    assert isinstance(train_rec_layer.cell.output_layers_net, TFNetwork)
    assert_equal(set(train_rec_layer.cell.output_layers_net.layers["output_prob"].params.keys()), {"b"})
  with make_scope() as session:
    print("Construct for search")
    search_net = TFNetwork(config=config, train_flag=False, eval_flag=True, search_flag=True)
    search_net.construct_from_dict(config.typed_dict["network"])


def test_name_scope():
  with make_scope() as session:
    n_in, n_out = 2, 3
    config = Config({"extern_data": {"data": {"dim": n_in}}})
    network = TFNetwork(config=config)
    net_dict = {
      "output": {
        "class": "linear", "n_out": n_out, "from": "data:data",
        "name_scope": "custom_name"},
    }
    network.construct_from_dict(net_dict)
    layer = network.layers["output"]
    param = layer.params["W"]
    assert param.name == "custom_name/W:0"


def test_name_scope_abs():
  with make_scope() as session:
    n_in, n_out = 2, 3
    config = Config({"extern_data": {"data": {"dim": n_in}}})
    network = TFNetwork(config=config)
    net_dict = {
      "output": {
        "class": "subnetwork",
        "subnetwork": {
          "layer1": {"class": "linear", "n_out": n_out, "from": "base:data:data"},
          "layer2": {"class": "linear", "n_out": n_out, "from": "base:data:data", "name_scope": "/custom_name"},
          "output": {"class": "combine", "kind": "add", "from": ["layer1", "layer2"]}
        }
      }
    }
    network.construct_from_dict(net_dict)
    layer1 = network.get_layer("output/layer1")
    layer2 = network.get_layer("output/layer2")
    param1 = layer1.params["W"]
    param2 = layer2.params["W"]
    assert param1.name == "output/layer1/W:0" and param2.name == "custom_name/W:0"


def test_name_scope_sub_empty():
  with make_scope() as session:
    n_in, n_out = 2, 3
    config = Config({"extern_data": {"data": {"dim": n_in}}})
    network = TFNetwork(config=config)
    net_dict = {
      "output": {
        "class": "subnetwork",
        "name_scope": "",
        "subnetwork": {
          "layer1": {"class": "linear", "n_out": n_out, "from": "base:data:data"},
          "layer2": {"class": "linear", "n_out": n_out, "from": "base:data:data", "name_scope": "layer1"},
          "output": {"class": "combine", "kind": "add", "from": ["layer1", "layer2"]}
        }
      }
    }
    network.construct_from_dict(net_dict)
    layer1 = network.get_layer("output/layer1")
    layer2 = network.get_layer("output/layer2")
    param1 = layer1.params["W"]
    param2 = layer2.params["W"]
    assert param1.name == "layer1/W:0" and param2.name == "layer1/W:0"


def test_name_scope_rec_sub_empty():
  with make_scope() as session:
    n_in, n_out = 2, 3
    config = Config({"extern_data": {"data": {"dim": n_in}}})
    network = TFNetwork(config=config)
    net_dict = {
      "output": {
        "class": "rec",
        "name_scope": "",
        "from": "data:data",
        "optimize_move_layers_out": False,
        "unit": {
          "layer1": {"class": "linear", "n_out": n_out, "from": "data:source", "is_output_layer": True},
          "layer2": {"class": "linear", "n_out": n_out, "from": "data:source", "is_output_layer": True,
                     "name_scope": "layer1"},
          "output": {"class": "combine", "kind": "add", "from": ["layer1", "layer2"]}
        }
      }
    }
    network.construct_from_dict(net_dict)
    layer1 = network.get_layer("output/layer1")
    layer2 = network.get_layer("output/layer2")
    param1 = layer1.params["W"]
    param2 = layer2.params["W"]
    assert param1.name == "layer1/W:0" and param2.name == "layer1/W:0"


def test_name_scope_share_params():
  with make_scope() as session:
    n_in, n_out = 3, 3
    config = Config({"extern_data": {"data": {"dim": n_in}}})
    network = TFNetwork(config=config)
    net_dict = {
      "layer1": {"class": "linear", "n_out": n_out, "from": "data:data"},
      "output": {"class": "linear", "n_out": n_out, "from": "layer1", "name_scope": "layer1"},
    }
    network.construct_from_dict(net_dict)
    l1 = network.layers["layer1"]
    l2 = network.layers["output"]
    assert_equal(set(l1.params.keys()), {"W", "b"})
    assert_equal(set(l2.params.keys()), {"W", "b"})
    assert l1.params["W"] is l2.params["W"]
    assert l1.params["b"] is l2.params["b"]
    assert_equal(set(network.get_trainable_params()), {l1.params["W"], l1.params["b"]})


def test_SliceLayer_output_placeholder():
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData())
    src = InternalLayer(name="src", network=net, output=Data(**{'name': 'src', "dim": 20, "sparse": True}))
    src.output.placeholder = tf.constant([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]], dtype=tf.int32)
    src.output.size_placeholder = {0: tf.constant([5, 3, 2], dtype=tf.int32)}
    layer = SliceLayer(
      name="slice", network=net, axis="T", slice_step=2, slice_start=1, sources=[src],
      output=SliceLayer.get_out_data_from_opts(
        name="slice", network=net, axis="T", slice_step=2, slice_start=1, sources=[src]))
    out, seq_lens = session.run([layer.output.placeholder, layer.output.size_placeholder[0]])
    print(out)
    print(seq_lens)
    assert isinstance(out, numpy.ndarray)
    assert isinstance(seq_lens, numpy.ndarray)
    assert_equal(
      out.tolist(),
      [[2, 4],
       [7, 9],
       [12, 14]])
    assert_equal(seq_lens.tolist(), [2, 1, 1])


def test_SliceLayer_NCHW():
  with make_scope() as session:
    import numpy as np
    net = TFNetwork(extern_data=ExternData())
    with tf_compat.v1.variable_scope("src_nchw"):
      src_nchw = InternalLayer(
        name="src_nchw", network=net,
        output=Data(**{
          "name": "src_nchw_output",
          "dim": 16,
          "shape": (16, None, 16),
          "batch_dim_axis": 0,
          "time_dim_axis": 2,
          "feature_dim_axis": 1,
          "sparse": False}))
      src_nchw.output.placeholder = tf_compat.v1.placeholder(shape=(None, 16, None, 16), dtype=tf.float32)
      src_nchw.output.size_placeholder = {1: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}
    with tf_compat.v1.variable_scope("src_nchw_feature_unspecified"):
      src_nchw_no_f = InternalLayer(
        name="src_nchw_feature_unspecified", network=net,
        output=Data(**{
          "name": "src_nchw_feature_unspecified_output",
          "dim": 16,
          "shape": (16, None, 16),
          "batch_dim_axis": 0,
          "time_dim_axis": 2,
          "feature_dim_axis": NotSpecified,
          "sparse": False}))
      src_nchw_no_f.output.placeholder = tf_compat.v1.placeholder(shape=(None, 16, None, 16), dtype=tf.float32)
      src_nchw_no_f.output.size_placeholder = {1: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}
    with tf_compat.v1.variable_scope("slice1"):
      slice1 = SliceLayer(
        name="slice1", network=net, axis="f", slice_step=2, sources=[src_nchw],
        output=SliceLayer.get_out_data_from_opts(name="slice1", axis="f", slice_step=2,
                                                 sources=[src_nchw]))
    with tf_compat.v1.variable_scope("slice2"):
      slice2 = SliceLayer(
        name="slice2", network=net, axis="f", slice_step=2, sources=[src_nchw_no_f],
        output=SliceLayer.get_out_data_from_opts(name="slice2", axis="f", slice_step=2,
                                                 sources=[src_nchw_no_f]))
    out1, out2 = session.run([slice1.output.placeholder, slice2.output.placeholder],
                             feed_dict={src_nchw.output.placeholder: np.random.rand(10, 16, 11, 16),
                                        src_nchw.output.size_placeholder[1]: np.full(shape=(10,), fill_value=11),
                                        src_nchw_no_f.output.placeholder: np.random.rand(10, 16, 11, 16),
                                        src_nchw_no_f.output.size_placeholder[1]: np.full(shape=(10,), fill_value=11)
                                        })
    assert out1.shape == (10, 8, 11, 16)
    assert slice1.output.dim == 8 and slice1.output.feature_dim_axis == 1
    assert out2.shape == (10, 16, 11, 8)
    assert slice2.output.dim == 8 and slice2.output.feature_dim_axis == 3


def test_pad_conv_slice():
  # https://github.com/rwth-i6/returnn/issues/1017
  with make_scope() as session:
    net = TFNetwork(config=Config({"extern_data": {"data": {"dim": 5}}}))
    net.construct_from_dict({
      "padding": {"class": "pad", "mode": "constant", "value": 0, "axes": ["T"], "padding": [(1, 1)], "from": "data"},
      "conv": {
        "class": "conv", "n_out": 5, "filter_size": (2,), "padding": "valid", "in_spatial_dims": ["T"],
        "from": "padding"},
      "output": {
        "class": "slice", "axis": "T", "slice_start": None, "slice_end": -1, "slice_step": None, "from": "conv"},
    })
    out = net.get_default_output_layer().output
    in_ = net.extern_data.get_default_input_data()
    assert_not_equal(in_.get_time_dim_tag(), out.get_time_dim_tag())
    net.initialize_params(session)
    session.run((out.placeholder, out.get_sequence_lengths()), feed_dict=make_feed_dict(net.extern_data))


def test_GatherLayer():
  with make_scope() as session:
    import numpy as np
    net = TFNetwork(extern_data=ExternData())
    batch_dim, gather_dim, time_dim, feature_dim1, feature_dim2 = 3, 4, 2, 1, 2
    # [B, D, T, F1]
    values = InternalLayer(
      name="values", network=net,
      output=Data(**{
        "name": "values",
        "batch_dim_axis": 0, "time_dim_axis": 2, "feature_dim_axis": 3,
        "shape": [gather_dim, None, feature_dim1], "sparse": False}))
    # [B, T, F2]
    position = InternalLayer(
      name="position", network=net,
      output=Data(**{
        "name": "position",
        "batch_dim_axis": 0, "time_dim_axis": 1, "shape": [None, feature_dim2],
        "sparse": True, "dim": gather_dim}))

    random = np.random.RandomState(42)
    values_seqs = random.rand(batch_dim, gather_dim, time_dim, feature_dim1).astype('float32')
    seq_lens = random.randint(1, time_dim, size=[batch_dim])
    seq_lens_tensor = tf.constant(seq_lens, dtype=tf.int32)
    values.output.placeholder = tf.constant(values_seqs, dtype=tf.float32)
    values.output.size_placeholder = {1: seq_lens_tensor}
    position_seqs = random.randint(low=0, high=gather_dim, size=[batch_dim, time_dim, feature_dim2])
    position.output.placeholder = tf.constant(position_seqs, dtype=tf.int32)
    position.output.size_placeholder = {0: seq_lens_tensor}
    position.output.sanity_check()
    values.output.sanity_check()

    # should become [B, T, F2, F1]
    layer = GatherLayer(
      name="gather", network=net,
      sources=[values], position=position, axis="dim:%i" % gather_dim,
      output=GatherLayer.get_out_data_from_opts(
        name="gather", sources=[values], position=position, axis="dim:%i" % gather_dim))
    layer.output.sanity_check()
    out_seqs, size = session.run([layer.output.placeholder, layer.output.size_placeholder.as_dict()])
    assert isinstance(out_seqs, numpy.ndarray)

    # test shapes
    print('shapes: values', values.output, 'position', position.output, 'output', layer.output)
    assert layer.output.batch_dim_axis == 0 and layer.output.time_dim_axis == 1
    assert layer.output.batch_shape == (None, None, feature_dim2, feature_dim1)
    assert np.shape(out_seqs) == (batch_dim, time_dim, feature_dim2, feature_dim1)
    assert layer.output.dtype == values.output.dtype
    assert np.array_equal(size[0], seq_lens)

    print('values [B, D, T, F1]:', values_seqs)
    print('position [B, T, F2]:', position_seqs)
    print('produced output [B, T, F2, F1]:', out_seqs)

    # test values
    for b in range(batch_dim):
      for t in range(seq_lens[b]):
        for f2 in range(feature_dim2):
          for f1 in range(feature_dim1):
            np.testing.assert_almost_equal(out_seqs[b, t, f2, f1], values_seqs[b, position_seqs[b, t, f2], t, f1])


def test_GatherLayer_constant_position():
  with make_scope() as session:
    import numpy as np
    net = TFNetwork(extern_data=ExternData())
    batch_dim, gather_dim, feature_dim1, feature_dim2 = 3, 4, 1, 2
    # [B, F1, D, F2]
    values = InternalLayer(
      name="values", network=net,
      output=Data(**{
        "name": "values",
        "batch_dim_axis": 0, "feature_dim_axis": 3, "time_dim_axis": None,
        "shape": [feature_dim1, gather_dim, feature_dim2]}))
    position = 3

    random = np.random.RandomState(42)
    values_seqs = random.rand(batch_dim, feature_dim1, gather_dim, feature_dim2).astype('float32')
    values.output.placeholder = tf.constant(values_seqs, dtype=tf.float32)
    values.output.sanity_check()

    # should become [B, F1, F2]
    layer = GatherLayer(
      name="gather", network=net,
      sources=[values], position=position, axis="dim:%i" % gather_dim,
      output=GatherLayer.get_out_data_from_opts(
        name="gather", sources=[values], position=position, axis="dim:%i" % gather_dim))
    layer.output.sanity_check()
    out_seqs = session.run(layer.output.placeholder)
    assert isinstance(out_seqs, numpy.ndarray)

    # test shapes
    print('shapes: values', values.output, 'position', position, 'output', layer.output)
    assert layer.output.batch_dim_axis == 0 and layer.output.feature_dim_axis == 2
    assert layer.output.batch_shape == (None, feature_dim1, feature_dim2)
    assert np.shape(out_seqs) == (batch_dim, feature_dim1, feature_dim2)
    assert layer.output.dtype == values.output.dtype

    print('values [B, F1, D, F2]:', values_seqs)
    print('position:', position)
    print('produced output [B, F1, F2]:', out_seqs)

    # test values
    for b in range(batch_dim):
      for f1 in range(feature_dim1):
        for f2 in range(feature_dim2):
          np.testing.assert_almost_equal(out_seqs[b, f1, f2], values_seqs[b, f1, position, f2])


def test_GatherLayer_search_beam():
  from returnn.tf.network import TFNetwork
  from returnn.config import Config
  with make_scope() as session:
    n_out = 5
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
        "classes": {"dim": n_out, "sparse": True}
      }})
    net = TFNetwork(config=config, search_flag=True)
    net.construct_from_dict({
      "output": {
        "class": "rec", "from": "data:data", "unit": {
          "position": {"class": "reinterpret_data", "from": "prev:output", "set_sparse": False},
          "gather": {"class": "gather", "from": "base:data:data", "position": "position", "axis": "t"},  # [B,T,slice,D]
          "prob": {"class": "softmax", "from": "gather", "target": "classes", "loss": "ce"},
          'output': {
            'class': 'choice', 'target': "classes", 'beam_size': 3, 'from': "prob", "input_type": "prob",
            "initial_output": 0}}}})


def test_GatherLayer_broadcast_dim():
  from returnn.tf.util.data import batch_dim
  head_dim = SpatialDim("head", 1)  # previously, this dim would match all others and therefore fail.
  round_dim = SpatialDim("round", 2)
  chunk_dim = SpatialDim("chunk")
  time_dim = SpatialDim("time")
  config = Config({"extern_data": {
    "source": {"dim_tags": [batch_dim, head_dim, time_dim]},
    "position": {"dim_tags": [batch_dim, head_dim, round_dim, chunk_dim], "dtype": "int32"}},
    "debug_print_layer_output_template": True})
  net = TFNetwork(config=config)
  net.construct_from_dict({
    "output": {
      'class': 'gather', 'from': 'data:source', 'position': 'data:position', 'axis': time_dim,
      'out_shape': {batch_dim, head_dim, round_dim, chunk_dim}}
  })


def test_GatherLayer_different_static_dims():
  # https://github.com/rwth-i6/returnn/issues/1219
  from returnn.tf.util.data import batch_dim
  head_dim = SpatialDim("head", 2)
  round_dim = SpatialDim("round", 2)
  chunk_dim = SpatialDim("chunk")
  time_dim = SpatialDim("time")
  config = Config({"extern_data": {
    "source": {"dim_tags": [batch_dim, head_dim, time_dim]},
    "position": {"dim_tags": [batch_dim, round_dim, chunk_dim], "dtype": "int32"}},
    "debug_print_layer_output_template": True})
  net = TFNetwork(config=config)
  net.construct_from_dict({
    "output": {
      'class': 'gather', 'from': 'data:source', 'position': 'data:position', 'axis': time_dim,
      'out_shape': {batch_dim, head_dim, round_dim, chunk_dim}}
  })


def test_SliceNdLayer():
  n_batch = 5
  n_time = 7
  n_dim = 11
  rnd = numpy.random.RandomState(42)
  seqs = rnd.randint(1, 100, (n_batch, n_time, n_dim)).astype("float32")  # all != 0
  seq_lens = numpy.array([n_time, n_time - 2, n_time - 3, n_time - 1, n_time - 4], dtype="int32")
  starts = numpy.array([2, 1, 3, n_time + 1, -1], dtype="int32")
  size = 5
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData())
    src = InternalLayer(name="src", network=net, output=Data(name="src", dim=n_dim))
    src.output.placeholder = tf.constant(seqs)
    src.output.size_placeholder = {0: tf.constant(seq_lens)}
    src.output.sanity_check()
    start = InternalLayer(
      name="start", network=net, output=Data(**{"name": "start", "dim": None, "sparse": True, "time_dim_axis": None}))
    start.output.placeholder = tf.constant(starts)
    start.output.sanity_check()
    kwargs = dict(name="slice", network=net, sources=[src], start=start, size=size)
    kwargs["output"] = SliceNdLayer.get_out_data_from_opts(**kwargs)
    layer = SliceNdLayer(**kwargs)
    print(layer)
    assert not layer.output.size_placeholder
    assert layer.output.batch_shape == (None, size, n_dim)
    out = session.run(layer.output.placeholder)
    print(out)
    assert isinstance(out, numpy.ndarray)
    assert out.shape == (n_batch, size, n_dim)
    for b in range(n_batch):
      s = starts[b]
      if s < 0:
        assert s + size > 0
        orig_seq = numpy.pad(seqs[b, :s + size], [(-s, 0), (0, 0)], "constant")
      else:
        orig_seq = seqs[b, s:s + size]
      if len(orig_seq) < size:
        orig_seq = numpy.pad(orig_seq, [(0, size - len(orig_seq)), (0, 0)], "constant")
      assert orig_seq.shape == (size, n_dim)
      orig_seq = numpy.where((numpy.arange(s, s + size) >= seq_lens[b])[:, None], 0.0, orig_seq)
      for t in range(size):
        numpy.testing.assert_equal(orig_seq[t], out[b, t])


def test_SliceNdLayer_dyn_size():
  n_batch = 4
  n_time = 7
  n_dim = 11
  rnd = numpy.random.RandomState(42)
  seqs = rnd.randint(1, 100, (n_batch, n_time, n_dim)).astype("float32")  # all != 0
  seq_lens = numpy.array([n_time, n_time - 2, n_time - 3, n_time - 1], dtype="int32")
  starts = numpy.array([2, 1, 3, n_time + 1], dtype="int32")
  size = None
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData())
    src = InternalLayer(name="src", network=net, output=Data(name="src", dim=n_dim))
    src.output.placeholder = tf.constant(seqs)
    src.output.size_placeholder = {0: tf.constant(seq_lens)}
    src.output.sanity_check()
    start = InternalLayer(
      name="start", network=net, output=Data(**{"name": "start", "dim": None, "sparse": True, "time_dim_axis": None}))
    start.output.placeholder = tf.constant(starts)
    start.output.sanity_check()
    kwargs = dict(name="slice", network=net, sources=[src], start=start, size=size)
    kwargs["output"] = SliceNdLayer.get_out_data_from_opts(**kwargs)
    layer = SliceNdLayer(**kwargs)
    print(layer)
    assert 0 in layer.output.size_placeholder
    assert layer.output.batch_shape == (None, size, n_dim)
    out = session.run(layer.output.placeholder)
    print(out)
    assert isinstance(out, numpy.ndarray)
    max_size = max(list(seq_lens - starts) + [0])
    assert out.shape == (n_batch, max_size, n_dim)
    for b in range(n_batch):
      s = starts[b]
      orig_seq = seqs[b, s:]
      if len(orig_seq) < max_size:
        orig_seq = numpy.pad(orig_seq, [(0, max_size - len(orig_seq)), (0, 0)], "constant")
      elif len(orig_seq) > max_size:
        orig_seq = orig_seq[:max_size]
      assert orig_seq.shape == (max_size, n_dim)
      orig_seq = numpy.where((numpy.arange(s, s + max_size) >= seq_lens[b])[:, None], 0.0, orig_seq)
      for t in range(max_size):
        numpy.testing.assert_equal(orig_seq[t], out[b, t])


def test_SliceNdLayer_multidimensional_start():
  with make_scope() as session:
    from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
    out_dim = FeatureDim("feat", 3)
    time_dim = SpatialDim("time")
    n_batch = 3
    max_seq_len = 10
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim_tags": [batch_dim, time_dim, out_dim]},
        "classes": {"dim_tags": [batch_dim, time_dim], "sparse_dim": out_dim}
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "output": {
        "class": "rec", "from": "data:data", "unit": {
          "start": {"class": "copy", "from": "prev:choice"},
          "slices": {"class": "slice_nd", "from": "base:data:data", "start": "start", "size": None},
          "output": {"class": "reduce", "from": "slices", "mode": "max", "axes": "stag:slice"},
          "prob": {"class": "softmax", "from": "data:source", "target": "classes", "loss": "ce"},
          'choice': {
            'class': 'choice', 'target': "classes", 'beam_size': 3, 'from': "prob", "input_type": "prob",
            "initial_output": 0,}}}})
    session.run(tf_compat.v1.global_variables_initializer())
    output_layer = net.layers["output"]
    starts = output_layer.cell.output_layers_net.layers["start"].output.get_placeholder_as_batch_major()
    segments = output_layer.cell.output_layers_net.layers["slices"].output.get_placeholder_as_batch_major()
    feed = make_feed_dict(net.extern_data.data.values(), n_batch=n_batch, n_time=max_seq_len, same_time=True)
    starts = session.run(starts, feed_dict=feed)
    segments = session.run(segments, feed_dict=feed)
    seq_lens = feed[net.extern_data.data["data"].size_placeholder[0]]
    input_data = feed[net.extern_data.data["data"].placeholder]
    max_size = numpy.amax(seq_lens[:, None] - starts)
    max_size = max(max_size, 0)
    assert segments.shape == (n_batch, max_seq_len, max_size, out_dim.dimension)
    for b in range(n_batch):
      for t in range(max_seq_len):
        s = starts[b, t]
        orig_seq = input_data[b, s:]
        if len(orig_seq) < max_size:
          orig_seq = numpy.pad(orig_seq, [(0, max_size - len(orig_seq)), (0, 0)], "constant")
        elif len(orig_seq) > max_size:
          orig_seq = orig_seq[:max_size]
        assert orig_seq.shape == (max_size, out_dim.dimension)
        orig_seq = numpy.where((numpy.arange(s, s + max_size) >= seq_lens[b])[:, None], 0.0, orig_seq)
        for t2 in range(max_size):
          numpy.testing.assert_equal(orig_seq[t2], segments[b, t, t2])


def test_SliceNdLayer_multidimensional_size():
  with make_scope() as session:
    from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
    out_dim = FeatureDim("feat", 3)
    time_dim = SpatialDim("time")
    n_batch = 3
    max_seq_len = 10
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim_tags": [batch_dim, time_dim, out_dim]},
        "classes": {"dim_tags": [batch_dim, time_dim], "sparse_dim": out_dim}
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "output": {
        "class": "rec", "from": "data:data", "unit": {
          "const1": {"class": "constant", "value": 1},
          "start": {"class": "reinterpret_data", "from": "prev:choice", "set_sparse": False},
          "size": {"class": "combine", "from": ["const1", "start"], "kind": "add"},
          "slices": {"class": "slice_nd", "from": "base:data:data", "start": "start", "size": "size"},
          "output": {"class": "reduce", "from": "slices", "mode": "max", "axes": "stag:slice"},
          "prob": {"class": "softmax", "from": "data:source", "target": "classes", "loss": "ce"},
          'choice': {
            'class': 'choice', 'target': "classes", 'beam_size': 3, 'from': "prob", "input_type": "prob",
            "initial_output": 0,}}}})
    session.run(tf_compat.v1.global_variables_initializer())
    output_layer = net.layers["output"]
    starts = output_layer.cell.output_layers_net.layers["start"].output.get_placeholder_as_batch_major()
    sizes = output_layer.cell.output_layers_net.layers["size"].output.get_placeholder_as_batch_major()
    segments = output_layer.cell.output_layers_net.layers["slices"].output.get_placeholder_as_batch_major()
    feed = make_feed_dict(net.extern_data.data.values(), n_batch=n_batch, n_time=max_seq_len, same_time=True)
    starts = session.run(starts, feed_dict=feed)
    sizes = session.run(sizes, feed_dict=feed)
    segments = session.run(segments, feed_dict=feed)
    seq_lens = feed[net.extern_data.data["data"].size_placeholder[0]]
    input_data = feed[net.extern_data.data["data"].placeholder]
    max_size = numpy.amax(sizes)
    max_size = max(max_size, 0)
    assert segments.shape == (n_batch, max_seq_len, max_size, out_dim.dimension)
    for b in range(n_batch):
      for t in range(max_seq_len):
        s = starts[b, t]
        size = sizes[b, t]
        end = min(s + size, seq_lens[b])
        orig_seq = input_data[b, s:end]
        if len(orig_seq) < max_size:
          orig_seq = numpy.pad(orig_seq, [(0, max_size - len(orig_seq)), (0, 0)], "constant")
        elif len(orig_seq) > max_size:
          orig_seq = orig_seq[:max_size]
        assert orig_seq.shape == (max_size, out_dim.dimension)
        orig_seq = numpy.where((numpy.arange(s, s + max_size) >= seq_lens[b])[:, None], 0.0, orig_seq)
        for t2 in range(max_size):
          numpy.testing.assert_equal(orig_seq[t2], segments[b, t, t2])


def test_SliceNdLayer_set_tag_on_size_tensor():
  with make_scope():
    n_out = 5
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
        "classes": {"dim": n_out, "sparse": True}
      }})
    net = TFNetwork(config=config, train_flag=True)
    # the construction of the "compare" layer will fail if set_tag_on_size_tensor is not called on the slice axis
    # inside of the SliceNdLayer
    net.construct_from_dict({
      "start": {"class": "range_in_axis", "from": "data", "axis": "b"},
      "slices": {"class": "slice_nd", "from": "data", "start": "start", "size": None},
      "output": {"class": "compare", "from": ["slices", "slices"], "kind": "equal"}
    })


def test_SliceNdLayer_start0():
  with make_scope() as session:
    from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
    time_dim = SpatialDim("time")
    feature_dim = FeatureDim("feature", 5)
    config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, feature_dim)}}})
    net = TFNetwork(config=config, train_flag=True)
    # the construction of the "compare" layer will fail if set_tag_on_size_tensor is not called on the slice axis
    # inside of the SliceNdLayer
    net.construct_from_dict({
      "downsample": {"class": "pool", "mode": "avg", "pool_size": 2, "from": "data"},
      "upsample": {"class": "resize", "axis": "T", "factor": 2, "from": "downsample"},
      "cutoff": {"class": "slice_nd", "from": "data", "size": time_dim},
      "output": {"class": "combine", "from": ["cutoff", "data"], "kind": "sub"}
    })
    session.run(net.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_SliceNdLayer_ReinterpretDataLayer():
  """
  https://github.com/rwth-i6/returnn/issues/851
  """
  from returnn.tf.util.data import DimensionTag
  new_slice_tag = DimensionTag(kind=DimensionTag.Types.Spatial, description="new-slice")
  with make_scope():
    n_out = 5
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
        "classes": {"dim": n_out, "sparse": True}
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "start": {"class": "reinterpret_data", "from": "data:classes", "set_sparse": False},
      "slices": {"class": "slice_nd", "from": "data", "start": "start", "size": None},
      "output": {
        "class": "reinterpret_data", "from": "slices", "set_dim_tags": {"stag:sliced-time:slices": new_slice_tag}}
    })


def test_WindowLayer_output_placeholder():
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData())
    src = InternalLayer(name="src", network=net, output=Data(name="src", dim=20, sparse=True))
    src.output.placeholder = tf.constant([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]], dtype=tf.int32)
    src.output.size_placeholder = {0: tf.constant([5, 3, 1], dtype=tf.int32)}
    layer = WindowLayer(
      name="window", network=net, axis="T", window_size=3, padding='valid', sources=[src],
      output=WindowLayer.get_out_data_from_opts(
        name="window", network=net, axis="T", window_size=3, padding='valid', sources=[src]))
    out, seq_lens = session.run([layer.output.placeholder, layer.output.size_placeholder[0]])
    print(out)
    print(seq_lens)
    assert isinstance(out, numpy.ndarray)
    assert isinstance(seq_lens, numpy.ndarray)
    assert_equal(
      out.tolist(),
      [[[1, 2, 3], [2, 3, 4], [3, 4, 5]],
       [[6, 7, 8], [7, 8, 9], [8, 9, 10]],
       [[11, 12, 13], [12, 13, 14], [13, 14, 15]]])
    assert_equal(seq_lens.tolist(), [3, 1, 0])


def test_conv_window_merge_dims():
  n_in = 1
  n_out = 13
  net_dict = {
    'conv_1': {'activation': 'abs',
               'class': 'conv',
               "from": "data:data",
               'filter_size': (4,),
               'n_out': 64,
               'padding': 'valid',
               'strides': 10},
    'pad_conv_1_time_dim': {'axes': 'time',
                            'class': 'pad',
                            'from': ['conv_1'],
                            'padding': 20},
    'conv_2': {'activation': 'abs',
               'class': 'conv',
               'filter_size': (2, 6),
               'from': ['pad_conv_1_time_dim'],
               'input_add_feature_dim': True,
               'n_out': 12,
               'padding': 'valid',
               'strides': 16},
    'flatten_conv': {'axes': ['dim:4', 'dim:12'],
                     'class': 'merge_dims',
                     'from': ['conv_2'],
                     'n_out': 48},
    'window_1': {'class': 'window',
                 'from': ['flatten_conv'],
                 'window_size': 17},
    'flatten_window': {'axes': ['dim:17', 'dim:48'],
                       'class': 'merge_dims',
                       'from': ['window_1']},
    'output': {'activation': None,
               'class': 'linear',
               'from': ['flatten_window'],
               'n_out': n_out},
  }
  config = Config({
    "num_outputs": n_out,
    "num_inputs": n_in,
    "debug_print_layer_output_template": True
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    print("extern data:")
    print(net.extern_data)
    # The construction itself is also the test.
    net.construct_from_dict(net_dict)
    out = net.get_default_output_layer()
    # Maybe this will not be the case in the future anymore;
    # however, if this test runs on CPU, currently the feature_dim_axis should always stay the default.
    # See also test_ConvLayer_feature_dim_unspecified.
    assert out.output.feature_dim_axis_or_unspecified in (NotSpecified, 1)


def test_ConvLayer_feature_dim_unspecified():
  n_in = 1
  n_out = 13
  net_dict = {
    'output': {'activation': 'abs',
               'class': 'conv',
               "from": "data:data",
               'filter_size': (4,),
               'n_out': 64,
               'padding': 'valid',
               'strides': 10}}
  config = Config({
    "num_outputs": n_out,
    "num_inputs": n_in,
    "debug_print_layer_output_template": True
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    print("extern data:")
    print(net.extern_data)
    net.construct_from_dict(net_dict)
    out = net.get_default_output_layer()
    # Maybe this will not be the case in the future anymore;
    # however, if this test runs on CPU, currently the feature_dim_axis should always stay the default.
    assert out.output.feature_dim_axis_or_unspecified is NotSpecified


def test_StftLayer():
  from returnn.tf.layers.signal_processing import StftLayer
  config = Config({"extern_data": {"data": {"dim": 1}}})
  with make_scope() as session:
    net = TFNetwork(config=config)
    in_layer = SourceLayer(name="input", network=net, data_key="data", output=net.extern_data.get_default_input_data())
    layer_desc = {
      "name": "stft",
      "network": net,
      "sources": [in_layer],
      "frame_size": 32,
      "frame_shift": 8,
    }
    stft_out = StftLayer.get_out_data_from_opts(**layer_desc)
    print("stft out:", stft_out)
    out_time = stft_out.get_time_dim_tag()
    assert in_layer.output.dim_tags[1].is_spatial_dim()
    assert out_time != in_layer.output.dim_tags[1]
    layer_desc["output"] = stft_out
    with tf_compat.v1.variable_scope("stft"):
      stft_layer = StftLayer(**layer_desc)
    net.layers["stft"] = stft_layer
    net.initialize_params(session)
    session.run(
      (stft_layer.output.placeholder, stft_layer.output.get_sequence_lengths()),
      feed_dict=make_feed_dict(net.extern_data, n_time=1024))


def test_IstftLayer():
  with make_scope() as session:
    frame_size = 32
    frame_shift = 8
    config = Config()
    config.update({
      "extern_data": {"data": {"dim": 1}},
      "network": {
        "stft": {"class": "stft", "frame_size": frame_size, "frame_shift": frame_shift, "from": "data"},
        "output": {"class": "istft", "frame_size": frame_size, "frame_shift": frame_shift, "from": "stft"},
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    out = network.get_default_output_layer().output.placeholder
    feed_dict = make_feed_dict(network.extern_data, n_time=1024)
    v_in = feed_dict[network.extern_data.data["data"].placeholder]
    v_out = session.run(out, feed_dict=feed_dict)
    # ignore border effects, see https://github.com/tensorflow/tensorflow/issues/16465
    numpy.testing.assert_allclose(v_in[:, frame_size:-frame_size, :], v_out[:, frame_size:-frame_size, :], rtol=1e-3)


def test_ConvLayer_time_dim_out():
  config = Config({"extern_data": {"data": {"dim": 7}}})
  with make_scope() as session:
    net = TFNetwork(config=config)
    in_layer = SourceLayer(name="input", network=net, data_key="data", output=net.extern_data.get_default_input_data())
    layer_desc = {
      'name': "conv",
      "network": net,
      "sources": [in_layer],
      'filter_size': (4,),
      'strides': 10,
      'padding': 'valid',
      'n_out': 64,
      'activation': 'abs'}
    conv_out = ConvLayer.get_out_data_from_opts(**layer_desc)
    print("conv out:", conv_out)
    out_time = conv_out.get_time_dim_tag()
    assert in_layer.output.dim_tags[1].is_spatial_dim()
    assert out_time != in_layer.output.dim_tags[1]
    layer_desc["output"] = conv_out
    with tf_compat.v1.variable_scope("conv"):
      conv_layer = ConvLayer(**layer_desc)
    net.layers["conv"] = conv_layer
    net.initialize_params(session)
    session.run(
      (conv_layer.output.placeholder, conv_layer.output.get_sequence_lengths()),
      feed_dict=make_feed_dict(net.extern_data))


def test_ConvLayer_get_out_data_from_opts_out_spatial_dims():
  from returnn.tf.util.data import batch_dim
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("input", 7)
  conv_dim = SpatialDim("conv")
  config = Config({"extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}}})
  with make_scope() as session:
    net = TFNetwork(config=config)
    layer_desc = {
      'name': "conv", "_name": "conv",
      "network": net, "_network": net,
      "from": "data",
      'filter_size': [4],
      'strides': 3,
      'padding': 'valid',
      'n_out': 13,
      "out_spatial_dims": [conv_dim],
    }
    ConvLayer.transform_config_dict(layer_desc, network=net, get_layer=net.get_layer)
    conv_out = ConvLayer.get_out_data_from_opts(**layer_desc)
    print("conv out:", conv_out)
    print("conv dim:", conv_dim, "time dim:", time_dim)
    assert conv_dim == ((-1) + time_dim + (-2)).ceildiv_right(3)
    assert conv_dim.get_same_base().derived_from_op
    with tf_compat.v1.variable_scope("conv"):
      conv_layer = ConvLayer(output=conv_out, **layer_desc)
    net.layers["conv"] = conv_layer
    net.initialize_params(session)
    session.run(
      (conv_layer.output.placeholder, conv_layer.output.get_sequence_lengths()),
      feed_dict=make_feed_dict(net.extern_data))


def test_ConvLayer_2d_device_based_opt():
  from returnn.tf.util.data import batch_dim
  for dev in ["cpu", "gpu"]:
    time_dim = SpatialDim("time")
    feat_dim = FeatureDim("input", 50)
    extra_dim = FeatureDim("extra", 1)
    config = Config({
      "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim, extra_dim]}},
      "device": dev,
    })
    with tf.Graph().as_default() as graph:
      # We don't use make_scope() or enter the session scope because this is consistent
      # to how it usually would be the case for RETURNN.
      net = TFNetwork(config=config)
      gpu_available = tf_util.is_gpu_available_in_session()
      print("GPU available:", gpu_available)
      assert gpu_available == (dev == "gpu")
      net_dict = {
        "output": {
          'class': 'conv',
          'from': 'data',
          'filter_size': (3, 3),
          'in_spatial_dims': ['T', 'dim:50'],
          'n_out': 32,
          'padding': 'same',
          'activation': None,
          'with_bias': True},
      }
      net.construct_from_dict(net_dict)
      conv_layer = net.get_default_output_layer()
      print("conv layer:", conv_layer)
      assert conv_layer.output.batch_dim_axis == 0
      if gpu_available:
        assert conv_layer.output.dim_tags[2:] == (time_dim, feat_dim)
        assert conv_layer.output.feature_dim_axis == 1 and conv_layer.output.dim == 32
      else:
        assert conv_layer.output.dim_tags[1:3] == (time_dim, feat_dim)
        assert conv_layer.output.feature_dim_axis == 3 and conv_layer.output.dim == 32


def test_ConvLayer_unrelated_dim():
  from returnn.tf.util.data import batch_dim
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("input", 7)
  other_dim = SpatialDim("other")
  config = Config({"extern_data": {"data": {"dim_tags": [batch_dim, time_dim, other_dim, feat_dim]}}})
  with make_scope() as session:
    net = TFNetwork(config=config)
    layer_desc = {
      'name': "conv", "_name": "conv",
      "network": net, "_network": net,
      "from": "data",
      'filter_size': [4],
      "in_spatial_dims": [time_dim],
      'strides': 3,
      'padding': 'valid',
      'n_out': 13,
    }
    ConvLayer.transform_config_dict(layer_desc, network=net, get_layer=net.get_layer)
    conv_out = ConvLayer.get_out_data_from_opts(**layer_desc)
    print("conv out:", conv_out)
    dyn_axes = conv_out.get_dynamic_axes()
    assert len(dyn_axes) == 2, "conv out: %r" % conv_out
    assert conv_out.get_axis_from_description(other_dim) in dyn_axes
    dyn_axes.remove(conv_out.get_axis_from_description(other_dim))
    out_spatial_dim = conv_out.dim_tags[dyn_axes[0]]
    assert out_spatial_dim not in net.extern_data.get_default_input_data().dim_tags
    assert conv_out.time_dim_axis == dyn_axes[0]
    with tf_compat.v1.variable_scope("conv"):
      conv_layer = ConvLayer(output=conv_out, **layer_desc)
    net.layers["conv"] = conv_layer
    net.initialize_params(session)
    session.run(
      (conv_layer.output.placeholder, conv_layer.output.get_sequence_lengths()),
      feed_dict=make_feed_dict(net.extern_data))


def test_conv_layer_NCHW():
  with make_scope() as session:
    import numpy as np
    net = TFNetwork(extern_data=ExternData())
    with tf_compat.v1.variable_scope("src_nhwc"):
      src_nhwc = InternalLayer(
        name="src_nhwc", network=net,
        output=Data(**{
          "name": "src_nhwc_output",
          "dim": 17,
          "shape": (None, 16, 17),
          "batch_dim_axis": 0,
          "time_dim_axis": 1,
          "feature_dim_axis": 3,
          "sparse": False}))
      src_nhwc.output.placeholder = tf_compat.v1.placeholder(shape=(None, None, 16, 17), dtype=tf.float32)
      src_nhwc.output.size_placeholder = {0: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}
    with tf_compat.v1.variable_scope("src_nchw"):
      src_nchw = InternalLayer(
        name="src_nchw", network=net,
        output=Data(**{
          "name": "src_nchw_output",
          "dim": 17,
          "shape": (17, None, 16),
          "batch_dim_axis": 0,
          "time_dim_axis": 2,
          "feature_dim_axis": 1,
          "sparse": False}))
      src_nchw.output.placeholder = tf_compat.v1.placeholder(shape=(None, 17, None, 16), dtype=tf.float32)
      src_nchw.output.size_placeholder = {1: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}

    filters = 64
    filter_size = (5, 5)
    strides = (1, 2)
    padding = "VALID"

    with tf_compat.v1.variable_scope("conv_nhwc_from_nhwc"):
      conv_nhwc_from_nhwc = ConvLayer(
        name="conv_nhwc_from_nhwc", network=net, n_out=filters, filter_size=filter_size,
        in_spatial_dims=["T", "dim:16"],
        padding=padding, strides=strides, auto_use_channel_first=False, sources=[src_nhwc],
        output=ConvLayer.get_out_data_from_opts(name="conv_nhwc_from_nhwc", n_out=filters,
                                                in_spatial_dims=["T", "dim:16"],
                                                filter_size=filter_size, padding=padding, strides=strides,
                                                auto_use_channel_first=False,
                                                network=net, sources=[src_nhwc]))
    with tf_compat.v1.variable_scope("conv_nchw_from_nhwc"):
      conv_nchw_from_nhwc = ConvLayer(
        name="conv_nchw_from_nhwc", network=net, n_out=filters, filter_size=filter_size,
        in_spatial_dims=["T", "dim:16"],
        padding=padding, strides=strides, auto_use_channel_first=True, sources=[src_nhwc],
        output=ConvLayer.get_out_data_from_opts(name="conv_nchw_from_nhwc", n_out=filters,
                                                in_spatial_dims=["T", "dim:16"],
                                                filter_size=filter_size, padding=padding, strides=strides,
                                                auto_use_channel_first=True,
                                                network=net, sources=[src_nhwc]))
    with tf_compat.v1.variable_scope("conv_nchw_from_nchw"):
      conv_nchw_from_nchw = ConvLayer(
        name="conv_nchw_from_nchw", network=net, n_out=filters, filter_size=filter_size,
        in_spatial_dims=["T", "dim:16"],
        padding=padding, strides=strides, auto_use_channel_first=True, sources=[src_nchw],
        output=ConvLayer.get_out_data_from_opts(name="conv_nchw_from_nchw", n_out=filters,
                                                in_spatial_dims=["T", "dim:16"],
                                                filter_size=filter_size, padding=padding, strides=strides,
                                                auto_use_channel_first=True,
                                                network=net, sources=[src_nchw]))
    tf_compat.v1.global_variables_initializer().run()
    out, seq_lens = session.run([conv_nhwc_from_nhwc.output.placeholder,
                                 conv_nhwc_from_nhwc.output.size_placeholder[0]],
                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 10, 16, 17),
                                           src_nhwc.output.size_placeholder[0]: np.full(shape=(10,), fill_value=10)}
                                )
    print(out.shape)
    assert_equal(out.shape, (10, 6, 6, 64))
    print(seq_lens)
    time_dim_axis = 1 if tf_util.is_gpu_available() else 0
    out, seq_lens = session.run([conv_nchw_from_nhwc.output.placeholder,
                                 conv_nchw_from_nhwc.output.size_placeholder[time_dim_axis]],
                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 10, 16, 17),
                                           src_nhwc.output.size_placeholder[0]: np.full(shape=(10,), fill_value=10)
                                })
    print(out.shape)
    if time_dim_axis == 1:
      assert_equal(out.shape, (10, 64, 6, 6))
    else:
      assert_equal(out.shape, (10, 6, 6, 64))
    print(seq_lens)
    if tf_util.is_gpu_available():
      out, seq_lens = session.run([conv_nchw_from_nchw.output.placeholder,
                                   conv_nchw_from_nchw.output.size_placeholder[1]],
                                  feed_dict={src_nchw.output.placeholder: np.random.rand(10, 17, 10, 16),
                                             src_nchw.output.size_placeholder[1]: np.full(shape=(10,), fill_value=10)
                                  })
      print(out.shape)
      assert_equal(out.shape, (10, 64, 6, 6))
      print(seq_lens)


def test_ConvLayer_empty_out():
  with make_scope() as session:
    net = TFNetwork(config=Config({"extern_data": {"data": {"dim": 5}}}))
    net.construct_from_dict({
      # Use filter_size 2 and T=1 to get 0 size out.
      # Using filter_size 3 would result in negative size according to the formula.
      # Actually I would have expected that TF also deals with this but this is not the case.
      "output": {"class": "conv", "n_out": 7, "filter_size": [2], "padding": "valid", "from": "data"},
    })
    out_ = net.layers["output"].output.copy_as_batch_spatial_major()
    print(out_)
    net.initialize_params(session)
    out, seq_lens = session.run(
      [out_.placeholder, out_.size_placeholder[0]],
      feed_dict=make_feed_dict(net.extern_data, n_time=1, n_batch=1))
    print(out)
    print(seq_lens)
    assert isinstance(out, numpy.ndarray)
    assert isinstance(seq_lens, numpy.ndarray)
    assert_equal(seq_lens.tolist(), [0])
    assert out.shape == (1, 0, 7)


def test_ConvLayer_stride_dilation():
  with make_scope() as session:
    net = TFNetwork(config=Config({"extern_data": {"data": {"dim": 5}}}))
    net.construct_from_dict({
      "output": {
        "class": "conv", "n_out": 7, "with_bias": False,
        "filter_size": (2,), "strides": (2,), "dilation_rate": (2,),
        "padding": "valid", "from": "data", "forward_weights_init": 1.0},
    })
    input = numpy.random.randint(low=0, high=10, size=(3, 20, 5))
    out_ = net.layers["output"].output.copy_as_batch_spatial_major()
    print(out_)
    net.initialize_params(session)
    out, seq_lens = session.run(
      [out_.placeholder, out_.size_placeholder[0]],
      feed_dict={
        net.extern_data.data["data"].size_placeholder[0]: numpy.array([20, 18, 17]),
        net.extern_data.data["data"].placeholder: input,
      }
    )
    print(out.shape)
    print(seq_lens)
    assert isinstance(out, numpy.ndarray)
    assert isinstance(seq_lens, numpy.ndarray)
    assert out.shape == (3, 9, 7)
    assert (seq_lens == [9, 8, 8]).all()


def test_pool_layer_NCHW():
  with make_scope() as session:
    import numpy as np
    net = TFNetwork(extern_data=ExternData())
    with tf_compat.v1.variable_scope("src_nhwc"):
      src_nhwc = InternalLayer(
        name="src_nhwc", network=net,
        output=Data(**{
          "name": "src_nhwc_output",
          "dim": 17,
          "shape": (None, 16, 17),
          "batch_dim_axis": 0,
          "time_dim_axis": 1,
          "feature_dim_axis": 3,
          "sparse": False}))
      src_nhwc.output.placeholder = tf_compat.v1.placeholder(shape=(None, None, 16, 17), dtype=tf.float32)
      src_nhwc.output.size_placeholder = {0: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}
    with tf_compat.v1.variable_scope("src_nchw"):
      src_nchw = InternalLayer(
        name="src_nchw", network=net,
        output=Data(**{
          "name": "src_nchw_output",
          "dim": 17,
          "shape": (17, None, 16),
          "batch_dim_axis": 0,
          "time_dim_axis": 2,
          "feature_dim_axis": 1,
          "sparse": False}))
      src_nchw.output.placeholder = tf_compat.v1.placeholder(shape=(None, 17, None, 16), dtype=tf.float32)
      src_nchw.output.size_placeholder = {1: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}

    pool_size = (5, 5)
    strides = (1, 2)
    padding = "VALID"

    with tf_compat.v1.variable_scope("pool_nhwc_from_nhwc"):
      pool_nhwc_from_nhwc = PoolLayer(
        name="pool_nhwc_from_nhwc", network=net, mode="max", pool_size=pool_size,
        in_spatial_dims=["T", "dim:16"],
        padding=padding, strides=strides, use_channel_first=False, sources=[src_nhwc],
        output=PoolLayer.get_out_data_from_opts(name="pool_nhwc_from_nhwc",
                                                pool_size=pool_size, padding=padding, strides=strides,
                                                in_spatial_dims=["T", "dim:16"],
                                                use_channel_first=False,
                                                network=net, sources=[src_nhwc]))
    with tf_compat.v1.variable_scope("pool_nchw_from_nhwc"):
      pool_nchw_from_nhwc = PoolLayer(
        name="pool_nchw_from_nhwc", network=net, mode="max", pool_size=pool_size,
        in_spatial_dims=["T", "dim:16"],
        padding=padding, strides=strides, use_channel_first=True, sources=[src_nhwc],
        output=PoolLayer.get_out_data_from_opts(name="pool_nchw_from_nhwc",
                                                pool_size=pool_size, padding=padding, strides=strides,
                                                in_spatial_dims=["T", "dim:16"],
                                                use_channel_first=True,
                                                network=net, sources=[src_nhwc]))
    with tf_compat.v1.variable_scope("pool_nchw_from_nchw"):
      pool_nchw_from_nchw = PoolLayer(
        name="pool_nchw_from_nchw", network=net, mode="max", pool_size=pool_size,
        in_spatial_dims=["T", "dim:16"],
        padding=padding, strides=strides, use_channel_first=True, sources=[src_nchw],
        output=PoolLayer.get_out_data_from_opts(name="pool_nchw_from_nchw",
                                                pool_size=pool_size, padding=padding, strides=strides,
                                                in_spatial_dims=["T", "dim:16"],
                                                use_channel_first=True,
                                                network=net, sources=[src_nchw]))
    tf_compat.v1.global_variables_initializer().run()
    out, seq_lens = session.run([pool_nhwc_from_nhwc.output.placeholder,
                                 pool_nhwc_from_nhwc.output.get_sequence_lengths()],
                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 11, 16, 17),
                                           src_nhwc.output.get_sequence_lengths(): np.full(shape=(10,), fill_value=11)}
                                )
    print(out.shape)
    assert_equal(out.shape, (10, 7, 6, 17))
    print(seq_lens)
    out, seq_lens = session.run([pool_nchw_from_nhwc.output.placeholder,
                                 pool_nchw_from_nhwc.output.get_sequence_lengths()],
                                feed_dict={src_nhwc.output.placeholder: np.random.rand(10, 11, 16, 17),
                                           src_nhwc.output.get_sequence_lengths(): np.full(shape=(10,), fill_value=11)
                                })
    print(pool_nchw_from_nhwc.output, out.shape)
    if pool_nchw_from_nhwc.output.feature_dim_axis == 1:
      assert_equal(out.shape, (10, 17, 7, 6))
    else:
      assert_equal(out.shape, (10, 7, 6, 17))
    print(seq_lens)
    if tf_util.is_gpu_available():
      out, seq_lens = session.run([pool_nchw_from_nchw.output.placeholder,
                                   pool_nchw_from_nchw.output.get_sequence_lengths()],
                                  feed_dict={src_nchw.output.placeholder: np.random.rand(10, 17, 11, 16),
                                             src_nchw.output.get_sequence_lengths(): np.full(shape=(10,), fill_value=11)
                                  })
      print(out.shape)
      assert_equal(out.shape, (10, 17, 7, 6))
      print(seq_lens)


def test_TransposedConvLayer_1d_time_major():
  # https://github.com/rwth-i6/returnn/issues/949
  from returnn.tf.util.data import batch_dim
  time_dim = SpatialDim("time")
  in_feat_dim = FeatureDim("in_feat", 5)
  config = Config({
    "extern_data": {
      "data": {"dim_tags": [time_dim, batch_dim, in_feat_dim]}
    }
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {
        "class": "transposed_conv",
        "activation": None,
        "filter_size": [3],
        "from": "data",
        "n_out": 7,
        "strides": [2],
      }})
    out = net.get_default_output_layer().output
    session.run(tf_compat.v1.global_variables_initializer())
    session.run(out.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_TransposedConvLayer_2d_simple():
  # https://github.com/rwth-i6/returnn/issues/595
  n_batch, n_time, n_in, n_out = 7, 3, 5, 13
  config = Config({
    "extern_data": {
      "data": {"dim": n_in, "shape": (n_in, None)}  # [B,D,T], i.e. batch-feature-major
    }
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    assert net.extern_data.get_default_input_data().is_batch_feature_major
    net.construct_from_dict({
      "Unflatten": {'class': 'split_dims', 'from': 'data', 'axis': 'T', 'dims': [-1, 1]},  # [B,D,T,1]
      "output": {  # [B,D',T,2]
        'class': 'transposed_conv', 'from': 'Unflatten',
        'activation': None, 'with_bias': True,
        "in_spatial_dims": ("T", "dim:1"),
        'n_out': n_out, 'filter_size': (1, 2), 'strides': (1, 1),
        'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0)},
    })
    out = net.get_default_output_layer().output.copy_as_batch_feature_major()
    assert out.batch_shape == (None, 13, None, 2)
    assert_equal(out.get_dim_tag(2), net.extern_data.get_default_input_data().get_time_dim_tag())
    assert out.dim_tags[1].dimension == n_out and out.dim_tags[3].dimension == 2
    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_in, n_time))
    session.run(tf_compat.v1.global_variables_initializer())
    out_v = session.run(out.placeholder, feed_dict={net.extern_data.data["data"].placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, n_out, n_time, 2)


def test_TransposedConvLayer_2d_2x2():
  n_batch, n_time, n_in, n_out = 7, 3, 5, 13
  config = Config({
    "extern_data": {
      "data": {"dim": n_in, "shape": (n_in, None)}  # [B,D,T], i.e. batch-feature-major
    }
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    assert net.extern_data.get_default_input_data().is_batch_feature_major
    net.construct_from_dict({
      "Unflatten": {'class': 'split_dims', 'from': 'data', 'axis': 'T', 'dims': [-1, 1]},  # [B,D,T,1]
      "output": {  # [B,D',T,2]
        'class': 'transposed_conv', 'from': 'Unflatten',
        'activation': None, 'with_bias': True,
        "in_spatial_dims": ("T", "dim:1"),
        'n_out': n_out, 'filter_size': (2, 2), 'strides': (2, 2),
        'padding': 'valid', 'output_padding': (0, 0), 'remove_padding': (0, 0)},
    })
    out = net.get_default_output_layer().output.copy_as_batch_feature_major()
    assert out.batch_shape == (None, 13, None, 2)
    assert_not_equal(out.get_dim_tag(2), net.extern_data.get_default_input_data().get_time_dim_tag())
    assert out.dim_tags[1].dimension == n_out and out.dim_tags[3].dimension == 2
    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_in, n_time))
    session.run(tf_compat.v1.global_variables_initializer())
    out_v = session.run(out.placeholder, feed_dict={net.extern_data.data["data"].placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, n_out, n_time * 2, 2)


def test_TransposedConvLayer_out_size_pool_pad_same():
  with make_scope() as session:
    from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
    time_dim = SpatialDim("time")
    feature_dim = FeatureDim("feature", 5)
    config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, feature_dim)}}})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "downsample": {
        "class": "pool", "mode": "avg", "pool_size": [2], "from": "data", "padding": "same"},  # ceildiv with same
      "upsample": {
        "class": "transposed_conv", "filter_size": [2], "from": "downsample", "padding": "same",
        "out_spatial_dims": [time_dim], "out_dim": feature_dim},
      "output": {"class": "combine", "from": ["upsample", "data"], "kind": "sub"}
    })
    net.initialize_params(session)
    session.run(net.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_TransposedConvLayer_out_size_pool_pad_valid():
  with make_scope() as session:
    from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
    time_dim = SpatialDim("time")
    feature_dim = FeatureDim("feature", 5)
    config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, feature_dim)}}})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "downsample": {
        "class": "pool", "mode": "avg", "pool_size": [2], "from": "data", "padding": "valid"},  # floordiv with valid
      "upsample": {
        "class": "transposed_conv", "filter_size": [2], "from": "downsample", "padding": "valid",
        "out_spatial_dims": [time_dim], "out_dim": feature_dim},
      "output": {"class": "combine", "from": ["upsample", "data"], "kind": "sub"}
    })
    net.initialize_params(session)
    session.run(net.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_ReduceLayer_NCHW():
  with make_scope() as session:
    import numpy as np
    net = TFNetwork(extern_data=ExternData())
    with tf_compat.v1.variable_scope("src_nchw"):
      src_nchw = InternalLayer(
        name="src_nchw", network=net,
        output=Data(**{
          "name": "src_nchw_output",
          "dim": 16,
          "shape": (16, None, 16),
          "batch_dim_axis": 0,
          "time_dim_axis": 2,
          "feature_dim_axis": 1,
          "sparse": False}))
      src_nchw.output.placeholder = tf_compat.v1.placeholder(shape=(None, 16, None, 16), dtype=tf.float32)
      src_nchw.output.size_placeholder = {1: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}
    with tf_compat.v1.variable_scope("reduce1"):
      reduce1 = ReduceLayer(
        name="reduce1", network=net, mode="max", axis="f", sources=[src_nchw],
        output=ReduceLayer.get_out_data_from_opts(name="reduce1", mode="max", axis="f",
                                                  sources=[src_nchw]))
    with tf_compat.v1.variable_scope("reduce2"):
      reduce2 = ReduceLayer(
        name="reduce2", network=net, mode="max", axis="b", sources=[src_nchw],
        output=ReduceLayer.get_out_data_from_opts(name="reduce2", mode="max", axis="b",
                                                  sources=[src_nchw]))
    out1, out2 = session.run([reduce1.output.placeholder, reduce2.output.placeholder],
                             feed_dict={src_nchw.output.placeholder: np.random.rand(10, 16, 11, 16),
                                        src_nchw.output.size_placeholder[1]: np.full(shape=(10,), fill_value=11)})
    assert_equal(out1.shape, (10, 11, 16))
    assert_equal(out2.shape, (16, 11, 16))
    assert reduce1.output.time_dim_axis == 1
    assert reduce2.output.feature_dim_axis == 0 and reduce2.output.dim == 16
    assert reduce2.output.batch_dim_axis is None


def test_Loss_NCHW():
  with make_scope() as session:
    import numpy as np
    net = TFNetwork(extern_data=ExternData())
    with tf_compat.v1.variable_scope("src_nchw"):
      src_nchw = InternalLayer(
        name="src_nchw", network=net,
        output=Data(**{
          "name": "src_nchw_output",
          "dim": 16,
          "shape": (16, None),
          "batch_dim_axis": 0,
          "time_dim_axis": 2,
          "feature_dim_axis": 1,
          "sparse": False}))
      src_nchw.output.placeholder = tf_compat.v1.placeholder(shape=(None, 16, None), dtype=tf.float32)
      src_nchw.output.size_placeholder = {1: tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)}

    with tf_compat.v1.variable_scope("activation"):
      activation = ActivationLayer(
        name="activation", activation="softmax", network=net, sources=[src_nchw],
        output=ActivationLayer.get_out_data_from_opts(name="activation", activation="softmax", network=net,
                                                      sources=[src_nchw]))

    target_placeholder = tf_compat.v1.placeholder(shape=(None, None, 16), dtype=tf.float32)
    target_size_placeholder = tf_compat.v1.placeholder(shape=(None,), dtype=tf.int32)
    target_data = Data(name="target", shape=(None, 16), placeholder=target_placeholder,
                       size_placeholder={0: target_size_placeholder},
                       time_dim_axis=1, feature_dim_axis=2)

    with tf_compat.v1.variable_scope("loss"):
      loss = CrossEntropyLoss(base_network=net)
      loss.init(output=activation.output, output_with_activation=activation.output_before_activation,
                target=target_data, layer=activation)

    random_input = np.random.rand(10, 16, 32)
    loss_out, out_flat = session.run([loss.get_value(), loss.output_before_softmax_flat],
                                     feed_dict={src_nchw.output.placeholder: random_input,
                                                src_nchw.output.size_placeholder[1]: np.full(shape=(10,), fill_value=32),
                                                target_placeholder: np.random.rand(10, 32, 16),
                                                target_size_placeholder: np.full(shape=(10,), fill_value=32)
                                                })
    print(loss_out)
    assert loss.output.feature_dim_axis == 2
    assert out_flat.shape == (320, 16)


def test_ResizeLayer_fill_value():
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData())
    src = InternalLayer(name="src", network=net, output=Data(name="src", dim=20, sparse=True))
    src.output.placeholder = tf.constant([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], dtype=tf.int32)
    src.output.size_placeholder = {0: tf.constant([5, 3], dtype=tf.int32)}
    layer = ResizeLayer(
      name="resize", network=net, factor=3, axis="T", kind="fill", fill_value=19, sources=[src],
      output=ResizeLayer.get_out_data_from_opts(
        name="resize", network=net, factor=3, axis="T", kind="fill", sources=[src]))
    out, seq_lens = session.run([layer.output.placeholder, layer.output.size_placeholder[0]])
    print(out)
    print(seq_lens)
    assert isinstance(out, numpy.ndarray)
    assert isinstance(seq_lens, numpy.ndarray)
    assert_equal(
      out.tolist(),
      [[1, 19, 19,  2, 19, 19,  3, 19, 19,  4, 19, 19,  5, 19, 19,],
       [6, 19, 19,  7, 19, 19,  8, 19, 19,  9, 19, 19, 10, 19, 19]])
    assert_equal(seq_lens.tolist(), [15, 9])


def test_ResizeLayer_fill_dropout():
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData())
    src = InternalLayer(name="src", network=net, output=Data(name="src", dim=20, sparse=True))
    src_seqs = [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]
    src_seq_lens = [5, 3]
    factor = 3
    fill_value = 19
    src.output.placeholder = tf.constant(src_seqs, dtype=tf.int32)
    src.output.size_placeholder = {0: tf.constant(src_seq_lens, dtype=tf.int32)}
    layer = ResizeLayer(
      name="resize", network=net,
      factor=factor, axis="T", kind="fill", fill_value=fill_value, fill_dropout=0.5, sources=[src],
      output=ResizeLayer.get_out_data_from_opts(
        name="resize", network=net, factor=factor, axis="T", kind="fill", sources=[src]))
    out, seq_lens = session.run([layer.output.placeholder, layer.output.size_placeholder[0]])
    print(out)
    print(seq_lens)
    assert isinstance(out, numpy.ndarray)
    assert isinstance(seq_lens, numpy.ndarray)
    # Non-deterministic output. But we can check some constraints.
    for i in range(len(src_seq_lens)):
      assert src_seq_lens[i] <= seq_lens[i] <= src_seq_lens[i] * factor
      assert_equal([s for s in out[i] if s != fill_value], src_seqs[i])


def test_ResizeLayer_BFT():
  n_batch, n_time, n_in = 7, 3, 5
  config = Config({
    "extern_data": {
      "data": {"dim": n_in, "shape": (n_in, None)}  # [B,F,T], i.e. batch-feature-major
    }
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    assert net.extern_data.get_default_input_data().is_batch_feature_major
    net.construct_from_dict({
      "output": {"class": "resize", "axis": "T", "factor": 2, "kind": "nn", "from": "data"}
    })
    out = net.get_default_output_layer().output.copy_as_batch_feature_major()
    in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_in, n_time))
    session.run(tf_compat.v1.global_variables_initializer())
    out_v = session.run(out.placeholder, feed_dict={net.extern_data.data["data"].placeholder: in_v})
    assert isinstance(out_v, numpy.ndarray)
    assert out_v.shape == (n_batch, n_in, n_time * 2)


def test_ResizeLayer_dynamic():
  n_batch, n_time, n_in = 2, 5, 3
  in_v = numpy.arange(0, n_batch * n_time * n_in).astype("float32").reshape((n_batch, n_time, n_in))
  in_seq_lens = numpy.array([5, 4])
  config = Config({
    "extern_data": {
      "data": {"shape": (None, n_in)},
      "factor": {"shape": (), "batch_dim_axis": None, "dtype": "float32"},
    }
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "output": {"class": "resize", "axis": "T", "factor": "data:factor", "kind": "nn", "from": "data"}
    })
    out = net.get_default_output_layer().output
    for factor in [0.5, 2.0]:
      out_v, out_lens = session.run(
        (out.placeholder, out.get_sequence_lengths()), feed_dict={
          net.extern_data.get_batch_info().dim: n_batch,
          net.extern_data.data["data"].placeholder: in_v,
          net.extern_data.data["data"].get_sequence_lengths(): in_seq_lens,
          net.extern_data.data["factor"].placeholder: factor,
        })
      assert isinstance(out_v, numpy.ndarray)
      assert out_v.shape == (n_batch, numpy.ceil(n_time * factor), n_in)
      numpy.testing.assert_equal(out_lens, numpy.ceil(in_seq_lens * factor).astype("int32"))


def test_PostfixInTimeLayer():
  with make_scope() as session:
    import numpy as np
    net = TFNetwork(extern_data=ExternData())
    batch = BatchInfo.make_global_batch_info(tf.constant(2))
    src = InternalLayer(name="src", network=net, output=Data(name="src", dim=2, dtype="int32", batch=batch))
    src_seqs = np.array([[[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]], [[6, 6], [7, 7], [8, 8], [0, 0], [0, 0]]])
    src_seq_lens = [5, 3]
    src.output.placeholder = tf.constant(src_seqs, dtype=tf.int32)
    src.output.size_placeholder = {0: tf.constant(src_seq_lens, dtype=tf.int32)}

    static_postfix = -7
    layer_postfix = InternalLayer(
      name="postfix", network=net, output=Data(name="postfix", dim=2, time_dim_axis=None, dtype="int32"))
    layer_postfix.output.placeholder = tf.constant([[-7, -8], [-9, -10]], dtype=tf.int32)
    layer_postfix_wo_batch = InternalLayer(
      name="postfix_wo_batch", network=net, output=Data(
        name="postfix_wo_batch", dim=2, time_dim_axis=None, batch_dim_axis=None, dtype="int32"))
    layer_postfix_wo_batch.output.placeholder = tf.constant([-7, -8], dtype=tf.int32)

    for postfix in [static_postfix, layer_postfix, layer_postfix_wo_batch]:
      for repeat in (1, 3):
        layer = PostfixInTimeLayer(
          name="postfix_in_time", network=net,
          sources=[src], postfix=postfix, repeat=repeat,
          output=PostfixInTimeLayer.get_out_data_from_opts(
            name="postfix_in_time", network=net, sources=[src], postfix=postfix, repeat=repeat))
        out, seq_lens = session.run([layer.output.placeholder, layer.output.size_placeholder[0]])
        print(out)
        print(seq_lens)
        assert isinstance(out, numpy.ndarray)
        assert isinstance(seq_lens, numpy.ndarray)
        assert out.shape == (2, 5 + repeat, 2)
        assert all(new_len == src_len + repeat for new_len, src_len in zip(seq_lens, src_seq_lens))
        assert out[0, src_seq_lens[0] - 1, 0] == src_seqs[0, src_seq_lens[0] - 1, 0]
        assert out[1, src_seq_lens[1] - 1, 0] == src_seqs[1, src_seq_lens[1] - 1, 0]
        assert out[0, src_seq_lens[0], 0] == -7
        assert out[0, src_seq_lens[0] + repeat - 1, 0] == -7


def test_DotLayer():
  with make_scope() as session:
    B = 2
    H = 3
    D = H * 5
    net = TFNetwork(extern_data=ExternData())
    a = InternalLayer(name="A", network=net, output=Data(name="A", shape=(None, H, D // H)))
    assert a.output.batch_dim_axis == 0
    assert a.output.time_dim_axis == 1
    assert a.output.shape == (None, H, D // H)
    assert a.output.dim == D // H
    a_seq_lens = [7, 3]
    assert len(a_seq_lens) == B
    a.output.placeholder = tf.reshape(
      tf.range(B * max(a_seq_lens) * D, dtype=tf.float32), (B, max(a_seq_lens), H, D // H))
    a.output.size_placeholder = {0: tf.constant(a_seq_lens, dtype=tf.int32)}
    b = InternalLayer(name="B", network=net, output=Data(name="B", shape=(H, D // H)))
    assert b.output.batch_dim_axis == 0
    assert b.output.shape == (H, D // H)
    assert b.output.dim == D // H
    b.output.placeholder = tf.reshape(tf.add(tf.range(B * D, dtype=tf.float32), 0.5), (B, H, D // H))
    kwargs = dict(
      name="dot", network=net, sources=[a, b], debug=True,
      red1="F", red2="F", var1="T", var2=None,
      add_var2_if_empty=False)
    layer = DotLayer(output=DotLayer.get_out_data_from_opts(**kwargs), **kwargs)
    print(layer, layer.output)
    assert layer.output.batch_dim_axis == 0
    assert layer.output.time_dim_axis == 2
    assert layer.output.shape == (H, None)
    out, seq_lens = session.run([layer.output.placeholder, layer.output.size_placeholder[1]])
    print(out)
    print(seq_lens)
    assert isinstance(out, numpy.ndarray)
    assert isinstance(seq_lens, numpy.ndarray)
    assert_equal(seq_lens.tolist(), a_seq_lens)
    assert_equal(out.shape, (B, H, max(a_seq_lens)))


def test_DotLayer2():
  """ Test if DotLayer can handle inputs which dont have a batch-dim
  """
  with make_scope() as session:
    B = 3
    S1, S2, R, V = 2, 4, 8, 16
    net = TFNetwork(extern_data=ExternData())

    a = InternalLayer(name='A',
                      network=net,
                      output=Data(name="A", shape=(S1, S2, R), batch_dim_axis=0, time_dim_axis=None))
    assert a.output.batch_dim_axis == 0
    assert a.output.time_dim_axis is None
    assert a.output.shape == (S1, S2, R)
    assert a.output.dim == R
    a.output.placeholder = tf.reshape(tf.range(B * S1 * S2 * R, dtype=tf.float32), (B, S1, S2, R))
    a.output.size_placeholder = {}

    b = InternalLayer(name='B',
                      network=net,
                      output=Data(name="B", shape=(S1, S2, R, V), batch_dim_axis=None, time_dim_axis=None))
    assert b.output.batch_dim_axis is None
    assert b.output.time_dim_axis is None
    assert b.output.shape == (S1, S2, R, V)
    assert b.output.dim == V
    b.output.placeholder = tf.reshape(tf.range(S1 * S2 * R * V, dtype=tf.float32), (S1, S2, R, V))
    b.output.size_placeholder = {}

    kwargs = dict(
      name="dot", network=net, sources=[a, b], debug=True,
      red1='F', red2='dim:%i' % R, var1='B', var2='F')
    layer = DotLayer(output=DotLayer.get_out_data_from_opts(**kwargs), **kwargs)
    print(layer, layer.output)
    assert layer.output.batch_dim_axis == 2
    assert layer.output.time_dim_axis is None
    assert layer.output.shape == (S1, S2, V)
    assert layer.output.batch_shape == (S1, S2, None, V)
    assert layer.output.dim == V
    out = session.run(layer.output.placeholder)
    assert isinstance(out, numpy.ndarray)
    assert_equal(out.shape, (S1, S2, B, V))


def test_DotLayer_linear_square_matrix():
  from returnn.tf.util.data import batch_dim
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("feature", dimension=3)
  config = Config({
    "extern_data": {
      "data": {"dim_tags": [batch_dim, time_dim, feat_dim]},
      "matrix_ambiguous": {"dim_tags": [feat_dim, feat_dim], "available_for_inference": True},
      "matrix_non_ambiguous": {
        "dim_tags": [feat_dim.copy(match_priority=1), feat_dim], "available_for_inference": True},
    },
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    try:
      net.construct_from_dict({
        "output": {
          "class": "dot", "from": ["data:data", "data:matrix_ambiguous"], "reduce": feat_dim
        },
      })
    except Exception as exc:
      print("Expected exception: %r" % exc)
      assert "must be unique" in str(exc)
    else:
      raise Exception("Expected exception but constructed layer: %s" % net.get_default_output_layer())
    net.construct_from_dict({
      "output": {
        "class": "dot", "from": ["data:data", "data:matrix_non_ambiguous"], "reduce": feat_dim
      },
    })
    out = net.get_default_output_layer().output
    assert out.dim_tags == (batch_dim, time_dim, feat_dim)
    session.run(out.placeholder, feed_dict=make_feed_dict(net.extern_data))


def test_DotLayer_mask_dyn_seq():
  batch = Dim(kind=Dim.Types.Batch, description="batch")
  time = SpatialDim("time")
  feat1 = FeatureDim("feature 1", dimension=3)
  feat2 = FeatureDim("feature 2", dimension=5)
  config = Config({
    "extern_data": {
      "src1": {"dim_tags": [batch, time, feat1]},
      "src2": {"dim_tags": [batch, time, feat2]},
    },
    "network": {
      "dot": {
        "class": "dot", "from": ["data:src1", "data:src2"], "is_output_layer": True,
        "red1": time, "red2": time, "var1": feat1, "var2": feat2
      },
    },
    "debug_print_layer_output_template": True,
  })

  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(config.typed_dict["network"])
    layer = net.layers["dot"]
    assert isinstance(layer, DotLayer)
    assert layer.output.dim_tags == (batch, feat1, feat2)
    assert layer._info_reduce_mask == "mask-source-1"

    feed_dict = make_feed_dict(net.extern_data)
    session.run(layer.output.placeholder, feed_dict=feed_dict)


def test_DotLayer_mask_dyn_seq_after_softmax():
  batch = Dim(kind=Dim.Types.Batch, description="batch")
  time = SpatialDim("time")
  feat1 = FeatureDim("feature 1", dimension=3)
  feat2 = FeatureDim("feature 2", dimension=5)
  config = Config({
    "extern_data": {
      "src1": {"dim_tags": [batch, time, feat1]},
      "src2": {"dim_tags": [batch, time, feat2]},
    },
    "network": {
      "sm1": {"class": "softmax_over_spatial", "from": "data:src1"},
      "dot": {
        "class": "dot", "from": ["sm1", "data:src2"], "is_output_layer": True,
        "red1": time, "red2": time, "var1": feat1, "var2": feat2
      },
    },
    "debug_print_layer_output_template": True,
  })

  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(config.typed_dict["network"])
    layer = net.layers["dot"]
    assert isinstance(layer, DotLayer)
    assert layer.output.dim_tags == (batch, feat1, feat2)
    assert layer._info_reduce_mask == "source-0-already-masked"

    feed_dict = make_feed_dict(net.extern_data)
    session.run(layer.output.placeholder, feed_dict=feed_dict)


def test_DotLayer_self_att_dyn_size_ext():
  batch_dim = Dim(kind=Dim.Types.Batch)
  heads_dim = SpatialDim("heads", dimension=8)
  classes_dim = Dim(kind=Dim.Types.Time, description="classes")
  keys_dim = Dim(
    kind=Dim.Types.Spatial, description="keys",
    dyn_size_ext=Data(name="keys_dyn_size", dim_tags=[classes_dim], dtype="int32", auto_create_placeholders=True))
  feature_dim = FeatureDim("feature", dimension=64)

  a = Data(name="att_weights", dim_tags=[batch_dim, heads_dim, classes_dim, keys_dim], auto_create_placeholders=True)
  b = Data(name="att_value", dim_tags=[keys_dim, batch_dim, heads_dim, feature_dim], auto_create_placeholders=True)
  print("a:", a)
  print("b:", b)

  config = Config({"debug_print_layer_output_template": True})
  config.update(dict(num_inputs=64 * 8, num_outputs=64 * 8))
  net = TFNetwork(config=config, train_flag=True, search_flag=False, eval_flag=False)

  a_lay = InternalLayer(name="att_weights", network=net, output=a)
  b_lay = InternalLayer(name="att_value", network=net, output=b)
  dot_kwargs = {"red1": keys_dim, "red2": keys_dim, "var1": classes_dim, "var2": feature_dim}
  dot = DotLayer.get_out_data_from_opts(name="dot", sources=[a_lay, b_lay], **dot_kwargs)
  DotLayer(network=net, name="dot", sources=[a_lay, b_lay], output=dot, **dot_kwargs)  # just check that it builds.


def test_DotLayer_sparse_input():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  classes_dim = FeatureDim("classes", dimension=5)
  embed_dim = FeatureDim("embed", dimension=3)
  config = Config({
    "extern_data": {
      "src": {"dim_tags": [batch_dim, time_dim], "sparse_dim": classes_dim},
      "embed": {"dim_tags": [classes_dim, embed_dim]},
    },
    "network": {
      "output": {"class": "dot", "from": ["data:src", "data:embed"], "reduce": classes_dim},
    },
    "debug_runtime_sanity_checks": True,
  })

  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(config.typed_dict["network"])
    layer = net.get_default_output_layer()
    assert layer.output.dim_tags == (batch_dim, time_dim, embed_dim)

    feed_dict = make_feed_dict(net.extern_data)
    session.run(layer.output.placeholder, feed_dict=feed_dict)


def test_DotLayer_dim_wrong_matching_same_dim_value():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("feat", dimension=5)
  feat2_dim = FeatureDim("other-feat", dimension=5)

  # First directly check DotLayer.get_out_data_from_opts.
  # This is more similar like we have it in returnn_common
  # and might trigger different errors due to the dim matching logic of DotLayer,
  # which behaves slightly different when there are no size_placeholders set yet,
  # see Dim.is_equal with unknown_spatial_matches.
  a = Data("a", dim_tags=[batch_dim, time_dim, feat_dim])
  b = Data("b", dim_tags=[batch_dim, time_dim, feat2_dim])
  net = TFNetwork(config=Config(), extern_data=ExternData())
  out = DotLayer.get_out_data_from_opts(
    name="dot",
    sources=[InternalLayer(name="a", network=net, output=a), InternalLayer(name="b", network=net, output=b)],
    reduce=time_dim)
  assert out.dim_tags == (batch_dim, feat_dim, feat2_dim)

  # Now full config.
  config = Config({
    "extern_data": {
      "a": {"dim_tags": [batch_dim, time_dim, feat_dim]},
      "b": {"dim_tags": [batch_dim, time_dim, feat2_dim]},
    },
    "network": {
      "output": {"class": "dot", "from": ["data:a", "data:b"], "reduce": time_dim},
    },
    "debug_runtime_sanity_checks": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(config.typed_dict["network"])
    layer = net.get_default_output_layer()
    assert layer.output.dim_tags == (batch_dim, feat_dim, feat2_dim)
    feed_dict = make_feed_dict(net.extern_data)
    session.run(layer.output.placeholder, feed_dict=feed_dict)


def test_DotLayer_dim_wrong_matching_derived():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  time_dim_2 = time_dim * 2
  assert time_dim_2.derived_from_tag == time_dim
  assert time_dim_2.get_same_derived_base() == time_dim
  feat_dim = FeatureDim("feat", dimension=5)

  # First directly check DotLayer.get_out_data_from_opts.
  # This is more similar like we have it in returnn_common
  # and might trigger different errors due to the dim matching logic of DotLayer,
  # which behaves slightly different when there are no size_placeholders set yet,
  # see Dim.is_equal with unknown_spatial_matches.
  a = Data("a", dim_tags=[batch_dim, time_dim, feat_dim])
  b = Data("b", dim_tags=[batch_dim, time_dim_2, feat_dim])
  net = TFNetwork(config=Config(), extern_data=ExternData())
  out = DotLayer.get_out_data_from_opts(
    name="dot",
    sources=[InternalLayer(name="a", network=net, output=a), InternalLayer(name="b", network=net, output=b)],
    reduce=feat_dim)
  assert out.dim_tags == (batch_dim, time_dim, time_dim_2)

  # Now full config.
  config = Config({
    "extern_data": {
      "a": {"dim_tags": [batch_dim, time_dim, feat_dim]},
      "b": {"dim_tags": [batch_dim, time_dim_2, feat_dim]},
    },
    "network": {
      "output": {"class": "dot", "from": ["data:a", "data:b"], "reduce": feat_dim},
    },
    "debug_runtime_sanity_checks": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(config.typed_dict["network"])
    layer = net.get_default_output_layer()
    assert layer.output.dim_tags == (batch_dim, time_dim, time_dim_2)
    feed_dict = make_feed_dict(net.extern_data)
    session.run(layer.output.placeholder, feed_dict=feed_dict)


def test_subnet_load_on_init():
  import tempfile
  model_tmp_dir = tempfile.mkdtemp("tmp-checkpoint")
  model_filename = model_tmp_dir + "/model"
  with make_scope() as session:
    config = Config()
    n_in, n_hidden, n_out = 2, 5, 3
    config.update({
      "num_outputs": n_out,
      "num_inputs": n_in,
      "network": {
        "l1": {"class": "linear", "activation": None, "n_out": n_hidden, "from": "data:data"},
        "output": {"class": "linear", "activation": None, "n_out": n_out, "from": ["l1"]}
      }
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    network.initialize_params(session)
    params_orig_dump = network.get_params_serialized(session)
    print("l1:")
    print(params_orig_dump.values_dict["l1"]["W"])
    print("output:")
    print(params_orig_dump.values_dict["output"]["W"])
    assert(params_orig_dump.values_dict["l1"]["W"].any())
    assert(params_orig_dump.values_dict["output"]["W"].any())
    network.save_params_to_file(filename=model_filename, session=session)

  with make_scope() as session:
    config = Config()
    config.update({
      "num_outputs": n_out,
      "num_inputs": n_in,
      "network": {
        "l0": {"class": "linear", "activation": None, "n_out": n_in, "from": "data:data"},
        "subnet": {"class": "subnetwork", "from": ["l0"], "load_on_init": model_filename, "subnetwork": {
          "l1": {"class": "linear", "activation": None, "n_out": n_hidden, "from": "data"},
          "output": {"class": "linear", "activation": None, "n_out": n_out, "from": ["l1"]}
        }},
        "output": {"class": "linear", "activation": None, "n_out": n_out, "from": ["subnet"]}
      }
    })
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    network.initialize_params(session)
    params_dump = network.get_params_serialized(session)
    params_dump_subnet = params_dump.values_dict["subnet"]
    for layer_name in ["l1", "output"]:
      layer_orig = params_orig_dump.values_dict[layer_name]
      for param_name in ["W", "b"]:
        param_orig = layer_orig[param_name]
        param_subnet = params_dump_subnet[layer_name + "/" + param_name]
        numpy.testing.assert_array_equal(param_orig, param_subnet)


def test_ReuseParams_rec():
  print("test_ReuseParams_rec()")
  numpy.set_printoptions(precision=15)
  num_inputs = 100
  num_outputs = 15
  config = Config()
  config.update({
    "num_inputs": num_inputs,
    "num_outputs": {"data": [num_inputs, 2], "classes": [num_outputs, 1]},  # dense output
    "network": {
      "out1":         {"class": "softmax", "from": ["rec_fwd"], "loss": "ce", "n_out": num_outputs},
      "out2":         {"class": "softmax", "from": ["rec_fwd_copy"], "loss": "ce", "n_out": num_outputs},
      "rec_fwd":      {"class": "rec", "direction": 1, "from": ["data"], "n_out": 300, "unit": "lstmp"},
      "rec_fwd_copy": {"class": "rec", "direction": 1, "from": ["data"], "n_out": 300, "unit": "lstmp", "reuse_params": "rec_fwd"}
    },
    "optimizer": {"class": "adam"},
    "target": "classes",
    "debug_grad_summaries": True,
    "debug_save_updater_vars": True,
    "debug_add_check_numerics_ops": True,
  })
  print("Creating network...")
  network = TFNetwork(config=config, train_flag=True)
  network.construct_from_dict(config.typed_dict["network"])
  random = numpy.random.RandomState(seed=1)
  def make_feed_dict(seq_len=10):
    return {
      network.extern_data.get_batch_info().dim: 1,
      network.extern_data.data["data"].placeholder: random.uniform(-1, 1, (1, seq_len, num_inputs)),
      network.extern_data.data["data"].size_placeholder[0]: numpy.array([seq_len]),
      network.extern_data.data["classes"].placeholder: random.randint(low=0, high=num_outputs, size=(1, seq_len)),
      network.extern_data.data["classes"].size_placeholder[0]: numpy.array([seq_len]),
    }
  print("Creating session...")
  with tf_compat.v1.Session() as session:
    print("Init params...")
    network.initialize_params(session=session)
    print("Testing reuse_params ...")
    feed = make_feed_dict(10)
    fwd_out, fwd_out_copy = session.run([network.layers["rec_fwd"].output.placeholder, network.layers["rec_fwd_copy"].output.placeholder], feed_dict=feed)
    numpy.testing.assert_array_equal(fwd_out, fwd_out_copy)


def test_ReuseParams_dep_loop():
  num_inputs = 10
  num_outputs = 15
  config = Config()
  config.update({
    "num_inputs": num_inputs,
    "num_outputs": {"data": [num_inputs, 2], "classes": [num_outputs, 1]},  # dense output
    "network": {
      "layer1": {"class": "rec", "from": "data", "n_out": 10, "unit": {
        "sub1": {
          "class": "linear", "from": ["data:source", "prev:output"], "activation": "relu", "n_out": 10
        },
        "sub2": {
          "class": "linear", "from": "sub1", "activation": "relu", "n_out": 10,
          "reuse_params": "base:layer2",  # circular dependency!
        },
        "output": {
          "class": "linear", "from": ["sub1", "sub2", "prev:output"], "activation": "relu", "n_out": 10
        }
      }},
      "layer2": {"class": "linear", "from": "layer1", "activation": "relu", "n_out": 10},
      "out": {"class": "softmax", "from": "layer2", "loss": "ce", "n_out": num_outputs},
    },
    "optimizer": {"class": "adam"},
    "target": "classes",
    "debug_print_layer_output_template": True,
  })
  print("Creating network...")
  network = TFNetwork(config=config, train_flag=True)
  network.construct_from_dict(config.typed_dict["network"])

  params = network.get_params_list()
  pprint(params)
  l1 = network.get_layer("layer1")
  from returnn.tf.layers.rec import RecLayer, _SubnetworkRecCell
  assert isinstance(l1, RecLayer)
  cell = l1.cell
  assert isinstance(cell, _SubnetworkRecCell)
  l1 = cell.net.layers["sub2"]
  assert isinstance(l1, LinearLayer)
  assert tf_util.has_control_flow_context(l1.output.placeholder)  # this should be in the loop
  l2 = network.get_layer("layer2")
  assert not tf_util.has_control_flow_context(l2.output.placeholder)  # outside the loop
  assert isinstance(l2, LinearLayer)
  assert set(l1.params.keys()) == set() and set(l2.params.keys()) == {"W", "b"}

  def make_feed_dict(seq_len=10):
    random = numpy.random.RandomState(seed=1)
    return {
      network.extern_data.get_batch_info().dim: 1,
      network.extern_data.data["data"].placeholder: random.uniform(-1, 1, (1, seq_len, num_inputs)),
      network.extern_data.data["data"].size_placeholder[0]: numpy.array([seq_len]),
      network.extern_data.data["classes"].placeholder: random.randint(low=0, high=num_outputs, size=(1, seq_len)),
      network.extern_data.data["classes"].size_placeholder[0]: numpy.array([seq_len]),
    }

  with tf_compat.v1.Session() as session:
    network.initialize_params(session=session)
    feed = make_feed_dict(10)
    # Not really needed (for testing reuse_params), but just test anyway.
    session.run(network.get_default_output_layer().output.placeholder, feed_dict=feed)


def test_ReuseParams_dep_loop_2():
  num_inputs = 10
  num_outputs = 15
  config = Config()
  config.update({
    "num_inputs": num_inputs,
    "num_outputs": {"data": [num_inputs, 2], "classes": [num_outputs, 1]},  # dense output
    "network": {
      "layer1": {"class": "rec", "from": "data", "n_out": 10, "unit": {
        "sub1": {
          "class": "linear", "from": ["data:source", "prev:output"], "activation": "relu", "n_out": 10
        },
        "sub2": {
          "class": "linear", "from": "sub1", "activation": "relu", "n_out": 10,
          "reuse_params": "base:layer2",  # circular dependency!
          "is_output_layer": True,
        },
        "output": {
          "class": "linear", "from": ["sub1", "sub2", "prev:output"], "activation": "relu", "n_out": 10
        }
      }},
      "layer2": {"class": "linear", "from": "layer1/sub2", "activation": "relu", "n_out": 10},
      "out": {"class": "softmax", "from": "layer2", "loss": "ce", "n_out": num_outputs},
    },
    "optimizer": {"class": "adam"},
    "target": "classes",
    "debug_print_layer_output_template": True,
  })
  print("Creating network...")
  network = TFNetwork(config=config, train_flag=True)
  network.construct_from_dict(config.typed_dict["network"])

  params = network.get_params_list()
  pprint(params)
  l1 = network.get_layer("layer1")
  from returnn.tf.layers.rec import RecLayer, _SubnetworkRecCell
  assert isinstance(l1, RecLayer)
  cell = l1.cell
  assert isinstance(cell, _SubnetworkRecCell)
  l1 = cell.net.layers["sub2"]
  assert isinstance(l1, LinearLayer)
  assert tf_util.has_control_flow_context(l1.output.placeholder)  # this should be in the loop
  l2 = network.get_layer("layer2")
  assert not tf_util.has_control_flow_context(l2.output.placeholder)  # outside the loop
  assert isinstance(l2, LinearLayer)
  assert set(l1.params.keys()) == set() and set(l2.params.keys()) == {"W", "b"}

  def make_feed_dict(seq_len=10):
    random = numpy.random.RandomState(seed=1)
    return {
      network.extern_data.get_batch_info().dim: 1,
      network.extern_data.data["data"].placeholder: random.uniform(-1, 1, (1, seq_len, num_inputs)),
      network.extern_data.data["data"].size_placeholder[0]: numpy.array([seq_len]),
      network.extern_data.data["classes"].placeholder: random.randint(low=0, high=num_outputs, size=(1, seq_len)),
      network.extern_data.data["classes"].size_placeholder[0]: numpy.array([seq_len]),
    }

  with tf_compat.v1.Session() as session:
    network.initialize_params(session=session)
    feed = make_feed_dict(10)
    # Not really needed (for testing reuse_params), but just test anyway.
    session.run(network.get_default_output_layer().output.placeholder, feed_dict=feed)


def test_ReuseParams_dep_loop_3():
  num_inputs = 10
  num_outputs = 15
  config = Config()
  config.update({
    "num_inputs": num_inputs,
    "num_outputs": {"data": [num_inputs, 2], "classes": [num_outputs, 1]},  # dense output
    "network": {
      "layer1": {"class": "rec", "from": "data", "n_out": 10, "unit": {
        "sub1": {
          "class": "linear", "from": ["data:source", "prev:output"], "activation": "relu", "n_out": 10,
          "is_output_layer": True,
        },
        "sub2": {
          "class": "linear", "from": "sub1", "activation": "relu", "n_out": 10,
          "reuse_params": "base:layer2",  # circular dependency!
        },
        "output": {
          "class": "linear", "from": ["sub1", "sub2", "prev:output"], "activation": "relu", "n_out": 10
        }
      }},
      "layer2": {"class": "linear", "from": "layer1/sub1", "activation": "relu", "n_out": 10},
      "out": {"class": "softmax", "from": "layer2", "loss": "ce", "n_out": num_outputs},
    },
    "optimizer": {"class": "adam"},
    "target": "classes",
    "debug_print_layer_output_template": True,
  })
  print("Creating network...")
  network = TFNetwork(config=config, train_flag=True)
  network.construct_from_dict(config.typed_dict["network"])

  params = network.get_params_list()
  pprint(params)
  l1 = network.get_layer("layer1")
  from returnn.tf.layers.rec import RecLayer, _SubnetworkRecCell
  assert isinstance(l1, RecLayer)
  cell = l1.cell
  assert isinstance(cell, _SubnetworkRecCell)
  l1 = cell.net.layers["sub2"]
  assert isinstance(l1, LinearLayer)
  assert tf_util.has_control_flow_context(l1.output.placeholder)  # this should be in the loop
  l2 = network.get_layer("layer2")
  assert not tf_util.has_control_flow_context(l2.output.placeholder)  # outside the loop
  assert isinstance(l2, LinearLayer)
  assert set(l1.params.keys()) == set() and set(l2.params.keys()) == {"W", "b"}

  def make_feed_dict(seq_len=10):
    random = numpy.random.RandomState(seed=1)
    return {
      network.extern_data.get_batch_info().dim: 1,
      network.extern_data.data["data"].placeholder: random.uniform(-1, 1, (1, seq_len, num_inputs)),
      network.extern_data.data["data"].size_placeholder[0]: numpy.array([seq_len]),
      network.extern_data.data["classes"].placeholder: random.randint(low=0, high=num_outputs, size=(1, seq_len)),
      network.extern_data.data["classes"].size_placeholder[0]: numpy.array([seq_len]),
    }

  with tf_compat.v1.Session() as session:
    network.initialize_params(session=session)
    feed = make_feed_dict(10)
    # Not really needed (for testing reuse_params), but just test anyway.
    session.run(network.get_default_output_layer().output.placeholder, feed_dict=feed)


def test_LossAsIs_custom_dim():
  config = Config()
  config.update({
    "extern_data": {
      "data": (40, 2),
      "classes": (10025, 1),
      "att_weights": {"shape": (None, None, 1)},
      "att_weights_sizes": {"shape": (None,), "dtype": "int32"}
    },
    "debug_print_layer_output_template": True,
  })
  print("Creating network...")
  network = TFNetwork(config=config, train_flag=True)
  net_dict = {
    "att_distill_loss": {
      "class": "eval", "from": ["energy", "att_weights"],
      "out_type": (lambda sources, **kwargs: sources[0].output.copy_template_excluding_spatial_dim(-1)),
      "eval": "softmax_cross_entropy_over_size(" +
              "logits=source(0, as_data=True, auto_convert=False)," +
              "labels=source(1, as_data=True, auto_convert=False))",
      "loss": "as_is"},
  }
  n_batch = 5
  n_enc_time = 11
  n_dec_time = 7
  with tf_compat.v1.Session() as session:
    enc_time = tf.constant([n_enc_time] * n_batch)
    dec_time = tf.constant([n_dec_time] * n_batch)
    network.add_layer(name="energy", layer_class=InternalLayer, output=Data(
      name="energy",
      shape=(None, None, 1), dim=1, batch_dim_axis=2,
      size_placeholder={0: dec_time, 1: enc_time},
      placeholder=tf.constant(numpy.random.normal(size=(n_dec_time, n_enc_time, n_batch, 1)).astype("float32"))))
    network.add_layer(name="att_weights", layer_class=InternalLayer, output=Data(
      name="att_weights",
      shape=(None, None, 1), dim=1, batch_dim_axis=0,
      size_placeholder={0: dec_time, 1: enc_time},
      placeholder=tf.expand_dims(
        tf.nn.softmax(
          tf.constant(numpy.random.normal(size=(n_batch, n_dec_time, n_enc_time)).astype("float32"))), -1)))
    network.construct_from_dict(net_dict)
    loss = session.run(network.get_total_loss())
    assert loss


def test_loss_cross_entropy():
  from returnn.tf.util.data import Dim, batch_dim, single_step_dim, SpatialDim, FeatureDim

  extern_data_data_dim_tags_1_time_dim = SpatialDim('time')
  extern_data_data_dim_tags_2_input_dim = FeatureDim('input', 3)
  extern_data_classes_sparse_dim_out_dim = FeatureDim('out', 5)

  extern_data_opts = {
    'data': {
      'dim_tags': (
        batch_dim,
        extern_data_data_dim_tags_1_time_dim,
        extern_data_data_dim_tags_2_input_dim
      ),
      'dtype': 'float32',
      'available_for_inference': True
    },
    'classes': {
      'dim_tags': (
        batch_dim,
        extern_data_data_dim_tags_1_time_dim
      ),
      'dtype': 'int32',
      'sparse_dim': extern_data_classes_sparse_dim_out_dim,
      'available_for_inference': True
    }
  }

  net_dict = {
    'dot': {
      'class': 'dot',
      'from': ['data:data', 'weight'],
      'reduce': extern_data_data_dim_tags_2_input_dim,
      'out_shape': {batch_dim, extern_data_data_dim_tags_1_time_dim, extern_data_classes_sparse_dim_out_dim}
    },
    'add': {
      'class': 'combine',
      'from': ['dot', 'bias'],
      'kind': 'add',
      'out_shape': {batch_dim, extern_data_data_dim_tags_1_time_dim, extern_data_classes_sparse_dim_out_dim}
    },
    'cross_entropy': {
      'class': 'subnetwork',
      'from': [],
      'subnetwork': {
        'sparse_softmax_cross_entropy_with_logits': {
          'class': 'sparse_softmax_cross_entropy_with_logits',
          'logits': 'base:add',
          'targets': 'base:data:classes',
          'axis': extern_data_classes_sparse_dim_out_dim,
          'out_shape': {batch_dim, extern_data_data_dim_tags_1_time_dim}
        },
        'output': {
          'class': 'copy',
          'from': 'sparse_softmax_cross_entropy_with_logits',
          'out_shape': {batch_dim, extern_data_data_dim_tags_1_time_dim}
        }
      },
      'out_shape': {batch_dim, extern_data_data_dim_tags_1_time_dim},
      'loss': 'as_is',
      'loss_scale': 1.0
    },
    'weight': {
      'class': 'variable',
      'shape': [
        extern_data_data_dim_tags_2_input_dim,
        extern_data_classes_sparse_dim_out_dim
      ]
    },
    'bias': {
      'class': 'variable',
      'shape': [
        extern_data_classes_sparse_dim_out_dim
      ]
    }
  }

  # First with train flag False. Only construct.
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData(extern_data_opts), train_flag=False)
    net.construct_from_dict(net_dict)

  # Now with train flag True. Try to calc loss.
  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData(extern_data_opts), train_flag=True)
    net.construct_from_dict(net_dict)
    loss = net.get_total_loss()
    print("loss tensor:", loss)
    assert isinstance(loss, tf.Tensor)
    tf_util.print_graph_output(loss, max_depth=4)
    from tensorflow.python.framework import tensor_util
    loss_const = tensor_util.constant_value(loss)
    assert loss_const is None
    net.initialize_params(session)
    session.run(loss, feed_dict=make_feed_dict(net.extern_data))


def test_loss_cross_entropy_as_is_optimize_flatten():
  from returnn.tf.util.data import Dim, batch_dim, single_step_dim, SpatialDim, FeatureDim

  extern_data_data_dim_tags_1_time_dim = SpatialDim('time')
  extern_data_data_dim_tags_2_input_dim = FeatureDim('input', 3)
  extern_data_classes_sparse_dim_out_dim = FeatureDim('out', 5)

  extern_data_opts = {
    'data': {
      'dim_tags': (
        batch_dim,
        extern_data_data_dim_tags_1_time_dim,
        extern_data_data_dim_tags_2_input_dim
      ),
      'dtype': 'float32',
      'available_for_inference': True
    },
    'classes': {
      'dim_tags': (
        batch_dim,
        extern_data_data_dim_tags_1_time_dim
      ),
      'dtype': 'int32',
      'sparse_dim': extern_data_classes_sparse_dim_out_dim,
      'available_for_inference': True
    }
  }

  net_dict = {
    'dot': {
      'class': 'dot',
      'from': ['data:data', 'weight'],
      'reduce': extern_data_data_dim_tags_2_input_dim,
      'out_shape': {batch_dim, extern_data_data_dim_tags_1_time_dim, extern_data_classes_sparse_dim_out_dim}
    },
    'add': {
      'class': 'combine',
      'from': ['dot', 'bias'],
      'kind': 'add',
      'out_shape': {batch_dim, extern_data_data_dim_tags_1_time_dim, extern_data_classes_sparse_dim_out_dim}
    },
    'cross_entropy': {
      'class': 'subnetwork',
      'from': [],
      'subnetwork': {
        'sparse_softmax_cross_entropy_with_logits': {
          'class': 'sparse_softmax_cross_entropy_with_logits',
          'logits': 'base:add',
          'targets': 'base:data:classes',
          'axis': extern_data_classes_sparse_dim_out_dim,
          'out_shape': {batch_dim, extern_data_data_dim_tags_1_time_dim}
        },
        'output': {
          'class': 'copy',
          'from': 'sparse_softmax_cross_entropy_with_logits',
          'out_shape': {batch_dim, extern_data_data_dim_tags_1_time_dim}
        }
      },
      'out_shape': {batch_dim, extern_data_data_dim_tags_1_time_dim},
      'loss': 'as_is',
      'loss_scale': 1.0
    },
    'weight': {
      'class': 'variable',
      'shape': [
        extern_data_data_dim_tags_2_input_dim,
        extern_data_classes_sparse_dim_out_dim
      ]
    },
    'bias': {
      'class': 'variable',
      'shape': [
        extern_data_classes_sparse_dim_out_dim
      ]
    }
  }

  with make_scope() as session:
    net = TFNetwork(extern_data=ExternData(extern_data_opts), train_flag=True)
    net.construct_from_dict(net_dict)
    loss = net.get_total_loss()
    print("loss tensor:", loss)
    assert isinstance(loss, tf.Tensor)
    tf_util.print_graph_output(loss, max_depth=4)

    assert loss.op.type == "Sum"  # reduce over flattened frames
    loss_ = loss.op.inputs[0]
    assert isinstance(loss_, tf.Tensor)
    assert loss_.op.type == "SparseSoftmaxCrossEntropyWithLogits"

    net.initialize_params(session)
    session.run(loss, feed_dict=make_feed_dict(net.extern_data))


def test_reduce_with_flatten():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  feature_dim = FeatureDim("feat", 1)
  config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, feature_dim), "dtype": "float32"}}})
  net_dict = {
    'exp': {
      'class': 'activation',
      'from': 'data:data',
      'activation': 'exp',
      'out_shape': {batch_dim, time_dim, feature_dim}
    },
    'mean_absolute_difference': {
      'class': 'subnetwork',
      'from': [],
      'subnetwork': {
        'sub': {
          'class': 'combine',
          'from': ['base:exp', 'base:data:data'],
          'kind': 'sub',
          'out_shape': {batch_dim, time_dim, feature_dim}
        },
        'abs': {
          'class': 'activation',
          'from': 'sub',
          'activation': 'abs',
          'out_shape': {batch_dim, time_dim, feature_dim}
        },
        'reduce': {
          'class': 'reduce',
          'from': 'abs',
          'mode': 'mean',
          'axis': feature_dim,
          'out_shape': {batch_dim, time_dim}
        },
        'output': {
          'class': 'copy',
          'from': 'reduce',
          'out_shape': {batch_dim, time_dim}
        }
      },
      'loss': 'as_is',
      'out_shape': {batch_dim, time_dim},
    },
}
  with make_scope():
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict(net_dict)
    net.get_total_loss()


def test_double_flatten_loss():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  _repeat_out_dim = time_dim * 5
  feature_dim = FeatureDim("feat", 1)
  config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, feature_dim), "dtype": "int32"}}})
  network = {
    'sub': {
      'class': 'combine',
      'from': ['data:data', 'data:data'],
      'kind': 'sub',
      'loss': 'as_is',
      'out_shape': {batch_dim, time_dim, feature_dim}
    },
    'repeat': {
      'class': 'repeat',
      'from': 'sub',
      'repetitions': 5,
      'axis': time_dim,
      'out_dim': _repeat_out_dim,
      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
    },
    'sub_0': {
      'class': 'combine',
      'from': ['repeat', 'repeat'],
      'kind': 'sub',
      'loss': 'as_is',
      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
    },
    'output': {
      'class': 'copy',
      'from': 'sub_0',
      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
    }
  }
  with make_scope():
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict(network)
    net.get_total_loss()


def test_double_flatten_loss_1079():
  # https://github.com/rwth-i6/returnn/issues/1079
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  _repeat_out_dim = time_dim * 5
  feature_dim = FeatureDim("feat", 1)
  config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, feature_dim), "dtype": "int32"}}})
  network = {
    'sub': {
      'class': 'combine',
      'from': ['data:data', 'data:data'],
      'kind': 'sub',
      'loss': 'as_is',
      'out_shape': {batch_dim, time_dim, feature_dim}
    },
    'repeat': {
      'class': 'repeat',
      'from': 'data:data',
      'repetitions': 5,
      'axis': time_dim,
      'out_dim': _repeat_out_dim,
      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
    },
    'sub_0': {
      'class': 'combine',
      'from': ['repeat', 'repeat'],
      'kind': 'sub',
      'loss': 'as_is',
      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
    },
    'output': {
      'class': 'copy',
      'from': 'sub_0',
      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
    }
  }
  with make_scope():
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict(network)
    net.get_total_loss()


def test_LossLayer_sublayers():
  from returnn.tf.util.basic import Dim
  n_in, n_out = 7, 11
  time_tag = SpatialDim("time")

  config = Config({
    "extern_data": {
      'data': {"dim": n_in, "same_dim_tags_as": {"t": time_tag}},
      'classes': {'dim': n_out, 'dtype': 'int64', 'sparse': True, "same_dim_tags_as": {"t": time_tag}},
      'prev-classes': {'dim': n_out, 'dtype': 'int64', 'sparse': True, "same_dim_tags_as": {"t": time_tag}}},
    "debug_print_layer_output_template": True,
  })
  net_dict = {
    'encoder-output': {"class": "linear", "activation": "relu", "n_out": 10, "from": "data:data"},

    'left-output': {'class': 'softmax', 'from': 'encoder-output', 'n_out': n_out},
    'left-output-ce': {'class': 'loss',
                       'from': 'left-output',
                       'loss': 'as_is',
                       'loss_': 'ce',
                       'loss_scale': 0,
                       'target_': 'prev-classes'},
    'left-err': {'class': 'copy', 'from': 'left-output-ce/error', 'loss': 'as_is', 'loss_scale': 0},
    'left-loss': {'class': 'copy', 'from': 'left-output-ce', 'loss': 'as_is'},

    'past-embed': {'activation': None, 'class': 'linear', 'from': ['data:prev-classes'], 'n_out': 10},
    'center-output': {'class': 'softmax', 'from': ['encoder-output', 'past-embed'], 'n_out': n_out},
    'center-output-ce': {'class': 'loss',
                         'from': 'center-output',
                         'loss': 'as_is',
                         'loss_': 'ce',
                         'loss_scale': 0,
                         'target_': 'classes'},
    'center-err': {'class': 'copy', 'from': 'center-output-ce/error', 'loss': 'as_is', 'loss_scale': 0},
    'center-loss': {'class': 'copy', 'from': 'center-output-ce', 'loss': 'as_is'},
  }
  print("Layers:", sorted(net_dict.keys()))

  print("Training")
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    # Check defaults available for inference.
    assert network.extern_data.data["data"].available_for_inference
    assert not network.extern_data.data["classes"].available_for_inference
    assert not network.extern_data.data["prev-classes"].available_for_inference
    network.construct_from_dict(net_dict)
    optimizer = tf_compat.v1.train.AdamOptimizer()
    fetches_dict = network.get_fetches_dict(should_train=True, should_eval=True, with_size=True)
    fetches_dict["optim_op"] = optimizer.minimize(network.get_objective())
    feed_dict = make_feed_dict(network.extern_data, same_time=True)
    session.run(tf_compat.v1.global_variables_initializer())
    for step in range(3):
      try:
        results = session.run(fetches_dict, feed_dict=feed_dict)
      except tf.errors.OpError as exc:
        help_on_tf_exception(
          session=session,
          exception=exc, fetches=fetches_dict, feed_dict=feed_dict,
          extern_data=network.extern_data)
        raise
      pprint(results)

  print("Forwarding")
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=False)
    network.construct_from_dict(net_dict)
    # Not sure if we would always expect that all layers with losses are constructed in this case,
    # but this is current behavior.
    # In any case, if we do that, they should infer the available_for_inference correctly.
    assert network.layers["encoder-output"].output.available_for_inference
    assert network.layers["left-output"].output.available_for_inference
    assert not network.layers["past-embed"].output.available_for_inference
    assert not network.layers["center-output"].output.available_for_inference
    assert not network.layers["center-output-ce"].output.available_for_inference
    assert not network.layers["center-loss"].output.available_for_inference
    assert not network.layers["center-err"].output.available_for_inference
    fetches_dict = network.get_fetches_dict(should_train=False, should_eval=False, with_size=True)
    feed_dict = make_feed_dict(network.extern_data, same_time=True)
    session.run(tf_compat.v1.global_variables_initializer())
    results = session.run(fetches_dict, feed_dict=feed_dict)
    pprint(results)

  print("Forwarding with fixed available-for-inference")
  with make_scope() as session:
    config.typed_dict["extern_data"]["prev-classes"]["available_for_inference"] = True
    network = TFNetwork(config=config, train_flag=False)
    network.construct_from_dict(net_dict)
    assert network.layers["encoder-output"].output.available_for_inference
    assert network.layers["left-output"].output.available_for_inference
    assert network.layers["past-embed"].output.available_for_inference
    assert network.layers["center-output"].output.available_for_inference
    assert not network.layers["center-output-ce"].output.available_for_inference
    assert not network.layers["center-loss"].output.available_for_inference
    assert not network.layers["center-err"].output.available_for_inference
    fetches_dict = network.get_fetches_dict(should_train=False, should_eval=False, with_size=True)
    feed_dict = make_feed_dict(network.extern_data, same_time=True)
    session.run(tf_compat.v1.global_variables_initializer())
    results = session.run(fetches_dict, feed_dict=feed_dict)
    pprint(results)


def test_EditDistanceLayer_greedy_ctc_decode():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  input_spatial_dim = SpatialDim("input-spatial")
  targets_spatial_dim = SpatialDim("targets-spatial")
  classes_dim = FeatureDim("classes", 10)
  blank_idx = classes_dim.dimension
  config = Config({
    "extern_data": {
      "logits": {"dim_tags": [batch_dim, input_spatial_dim, classes_dim + 1]},
      "targets": {"dim_tags": [batch_dim, targets_spatial_dim], "sparse_dim": classes_dim},
    },
  })
  net_dict = {
    "argmax": {"class": "reduce", "from": "data:logits", "axis": "F", "mode": "argmax"},
    "ctc_decode": {"class": "subnetwork", "from": "argmax", "subnetwork": {
      # tf_util.sparse_labels_with_seq_lens
      "shift_right": {
        "class": "shift_axis", "from": "data", "axis": "T", "amount": 1, "pad_value": -1, "adjust_size_info": False},
      "unique_mask": {
        "class": "compare", "from": ["data", "shift_right"], "kind": "not_equal"},
      "non_blank_mask": {
        "class": "compare", "from": "data", "kind": "not_equal", "value": blank_idx},
      "mask": {"class": "combine", "kind": "logical_and", "from": ["unique_mask", "non_blank_mask"]},
      "output": {
        "class": "masked_computation", "from": "data", "mask": "mask",
        "unit": {"class": "copy", "from": "data"}},
    }},
    "edit_dist": {"class": "edit_distance", "a": "ctc_decode", "b": "data:targets", "is_output_layer": True},
  }

  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)

    logits = network.extern_data.data["logits"]
    targets = network.extern_data.data["targets"]
    decoded_sparse = tf_util.ctc_greedy_decode(
      logits=logits.get_placeholder_as_time_major(), seq_lens=logits.get_sequence_lengths(), time_major=True)
    decoded_dense = tf_compat.v1.sparse_to_dense(
      sparse_indices=decoded_sparse.indices,
      sparse_values=decoded_sparse.values,
      output_shape=decoded_sparse.dense_shape)
    targets_sparse = tf_util.sparse_labels(targets.placeholder, seq_lens=targets.get_sequence_lengths())
    error = tf.edit_distance(
      hypothesis=tf.cast(decoded_sparse, targets_sparse.dtype), truth=targets_sparse, normalize=False)

    argmax = network.layers["argmax"].output
    decoded_net = network.layers["ctc_decode"].output
    print("decoded out:", decoded_net)
    error_net = network.layers["edit_dist"].output
    print("error out:", error_net)

    n_batch = 5
    n_logits_time = 20
    n_targets_time = 7
    rnd = numpy.random.RandomState(42)
    targets_np = rnd.randint(low=0, high=classes_dim.dimension, size=(n_batch, n_targets_time))
    targets_sizes = [7, 6, 5, 4, 3]
    logits_np = rnd.uniform(size=(n_batch, n_logits_time, classes_dim.dimension + 1)).astype("float32")
    for b in range(n_batch):
      # Random correct.
      for _ in range(rnd.randint(1, 7)):
        size = rnd.randint(1, 6)
        t = rnd.randint(0, n_logits_time - size + 1)
        logits_np[b, t:t + size, targets_np[b, rnd.randint(0, targets_sizes[b])]] = rnd.uniform(1., 2.)
      # Random repeats.
      for _ in range(rnd.randint(1, 5)):
        size = rnd.randint(1, 6)
        t = rnd.randint(0, n_logits_time - size + 1)
        logits_np[b, t:t + size] = logits_np[b, t:t+1]
      # Random blanks.
      for _ in range(rnd.randint(1, 5)):
        size = rnd.randint(1, 6)
        t = rnd.randint(0, n_logits_time - size + 1)
        logits_np[b, t:t + size, blank_idx] = rnd.uniform(1., 2.)
    logits_sizes = [20, 19, 18, 17, 16]
    feed_dict = {
      logits.placeholder: logits_np,
      logits.get_sequence_lengths(): logits_sizes,
      targets.placeholder: targets_np,
      targets.get_sequence_lengths(): targets_sizes,
    }
    argmax_np, dec1_np, err1_np, dec2_np, dec2_sizes, err2_np = session.run(
      (argmax.placeholder,
       decoded_dense, error,
       decoded_net.get_placeholder_as_batch_major(), decoded_net.get_sequence_lengths(),
       error_net.placeholder),
      feed_dict=feed_dict)
    print("targets:", targets_np)
    print("targets sizes:", targets_sizes)
    print("argmax:", argmax_np)
    print("decoded ref:", dec1_np)
    print("error ref:", err1_np)
    print("decoded:", dec2_np)
    print("decoded sizes:", dec2_sizes)
    print("error:", err2_np)
    numpy.testing.assert_array_equal(dec1_np, dec2_np)
    numpy.testing.assert_array_equal(err1_np, err2_np)


def test_param_variational_noise():
  from returnn.tf.util.basic import print_graph_output, find_ops_with_tensor_input
  config = Config({
    "debug_print_layer_output_template": True,
    "param_variational_noise": 0.075,
    "extern_data": {"data": {"dim": 7}}
  })
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    # Do subnetwork by intention, to test when we have multiple variable scopes.
    network.construct_from_dict({
      "output": {
        "class": "subnetwork",
        "from": "data:data",
        "subnetwork": {
          "output": {"class": "linear", "n_out": 13, "activation": "tanh", "from": "data"}
        }
      }
    })
    out = network.get_default_output_layer().output.placeholder
    print("output:")
    print_graph_output(out)
    params = network.get_params_list()
    print("params:", params)
    assert len(params) == 2  # weights and bias
    for param in params:
      print("param:", param)
      ops = find_ops_with_tensor_input(param, fetches=out)
      print("param graph:")
      print_graph_output(ops)
      assert len(ops) == 1 and "_variational_noise/" in ops[0].name


def test_LinearLayer_simple_train():
  config = Config()
  n_in, n_out = 7, 3
  config.update({
    "extern_data": {
      "data": (n_in, 2),
      "classes": (n_out, 1),
    },
    "debug_print_layer_output_template": True,
  })
  print("Creating network...")
  with tf.Graph().as_default():
    network = TFNetwork(config=config, train_flag=True)

    net_dict = {}
    layer_n_out = 10
    layer_common_args = {"class": "linear", "activation": "relu", "n_out": layer_n_out, "L2": 0.01}

    def layer(sources, **kwargs):
      args = kwargs.copy()
      for k, v in layer_common_args.items():
        args.setdefault(k, v)
      args.setdefault("from", sources)
      return args

    def make_network(num_layers):
      sources = ["data"]
      for i in range(num_layers):
        net_dict["layer%i" % i] = layer(sources=sources)
        sources = ["layer%i" % i]
      net_dict["output"] = {"class": "softmax", "loss": "ce", "from": sources}

    make_network(num_layers=3)
    network.construct_from_dict(net_dict)
    data_input = network.extern_data.get_default_input_data()
    data_target = network.extern_data.get_default_target_data()
    optimizer = tf_compat.v1.train.AdamOptimizer()
    network.maybe_construct_objective()
    update_op = optimizer.minimize(network.get_objective())
    n_batch = 5
    n_time = 11
    rnd = numpy.random.RandomState(42)
    with tf_compat.v1.Session() as session:
      session.run(tf_compat.v1.global_variables_initializer())
      for step in range(5):
        info, _ = session.run(
          (network.get_fetches_dict(), update_op),
          feed_dict={
            data_input.placeholder: rnd.normal(size=(n_batch, n_time, n_in)).astype("float32"),
            data_input.size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
            data_target.placeholder: rnd.randint(0, n_out, size=(n_batch, n_time), dtype="int32"),
            data_target.size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
          })
        print("step:", step, "info:", info)


def test_flat_net_construction():
  config = Config()
  n_in, n_out = 7, 3
  config.update({
    "extern_data": {
      "data": (n_in, 2),
      "classes": (n_out, 1),
    },
    "flat_net_construction": True,
    "debug_print_layer_output_template": True,
  })
  print("Creating network...")
  with tf.Graph().as_default():
    network = TFNetwork(config=config, train_flag=True)

    net_dict = {
      "pre0": {"class": "linear", "activation": "tanh", "from": "data", "n_out": 10},
      "pre1": {"class": "linear", "activation": "tanh", "from": "pre0", "n_out": 10},
      "pre2": {"class": "linear", "activation": "tanh", "from": "pre1", "n_out": 10}
    }
    layer_common_args = {"class": "copy"}

    def layer(sources, **kwargs):
      args = kwargs.copy()
      for k, v in layer_common_args.items():
        args.setdefault(k, v)
      args.setdefault("from", sources)
      return args

    def make_network(num_layers):
      sources = ["pre2"]
      for i in range(num_layers):
        net_dict["layer%i" % i] = layer(sources=sources)
        sources = ["layer%i" % i]
      net_dict["output"] = {"class": "softmax", "loss": "ce", "from": sources}

    make_network(num_layers=5000)
    network.construct_from_dict(net_dict)
    data_input = network.extern_data.get_default_input_data()
    data_target = network.extern_data.get_default_target_data()
    optimizer = tf_compat.v1.train.AdamOptimizer()
    network.maybe_construct_objective()
    update_op = optimizer.minimize(network.get_objective())
    n_batch = 5
    n_time = 11
    rnd = numpy.random.RandomState(42)
    with tf_compat.v1.Session() as session:
      session.run(tf_compat.v1.global_variables_initializer())
      for step in range(5):
        info, _ = session.run(
          (network.get_fetches_dict(), update_op),
          feed_dict={
            data_input.placeholder: rnd.normal(size=(n_batch, n_time, n_in)).astype("float32"),
            data_input.size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
            data_target.placeholder: rnd.randint(0, n_out, size=(n_batch, n_time), dtype="int32"),
            data_target.size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
          })
        print("step:", step, "info:", info)


def test_SyntheticGradientLayer():
  """
  Tests :class:`SyntheticGradientLayer`.
  """
  config = Config()
  n_in, n_out = 7, 3
  config.update({
    "extern_data": {
      "data": (n_in, 2),
      "classes": (n_out, 1),
    },
    "debug_print_layer_output_template": True,
  })
  print("Creating network...")
  with tf.Graph().as_default():
    network = TFNetwork(config=config, train_flag=True)

    net_dict = {}
    layer_n_out = 10
    layer_common_args = {"class": "linear", "activation": "relu", "n_out": layer_n_out, "L2": 0.01}

    def layer(sources, **kwargs):
      args = kwargs.copy()
      for k, v in layer_common_args.items():
        args.setdefault(k, v)
      args.setdefault("from", sources)
      return args

    def make_network(num_layers):
      sources = ["data"]
      for i in range(num_layers):
        net_dict["layer%i" % i] = layer(sources=sources)
        sources = ["layer%i" % i]
        net_dict["predict_grad%i" % i] = layer(sources=sources)
        net_dict["syn_grad%i" % i] = {"class": "synthetic_gradient", "gradient": "predict_grad%i" % i, "from": sources}
        sources = ["syn_grad%i" % i]
      net_dict["output"] = {"class": "softmax", "loss": "ce", "from": sources}

    make_network(num_layers=3)
    network.construct_from_dict(net_dict)
    data_input = network.extern_data.get_default_input_data()
    data_target = network.extern_data.get_default_target_data()
    from returnn.tf.updater import Updater
    updater = Updater(config=config, network=network, initial_learning_rate=0.001)
    updater.set_trainable_vars(tf_compat.v1.trainable_variables())
    update_op = updater.get_optim_op()
    assert updater.optim_meta_losses_dict
    fetches = network.get_fetches_dict()
    fetches.update(updater.optim_meta_losses_dict)

    n_batch = 5
    n_time = 11
    rnd = numpy.random.RandomState(42)
    with tf_compat.v1.Session() as session:
      session.run(tf_compat.v1.variables_initializer(tf_compat.v1.global_variables() + [network.global_train_step_var]))
      for step in range(5):
        info, _ = session.run(
          (fetches, update_op),
          feed_dict={
            data_input.placeholder: rnd.normal(size=(n_batch, n_time, n_in)).astype("float32"),
            data_input.size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
            data_target.placeholder: rnd.randint(0, n_out, size=(n_batch, n_time), dtype="int32"),
            data_target.size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
          })
        print("step:", step, "info:", info)


def test_TikhonovRegularizationLayer():
  """
  Tests :class:`TikhonovRegularizationLayer`.
  """
  config = Config()
  n_in, n_out = 7, 3
  config.update({
    "extern_data": {
      "data": (n_in, 2),
      "classes": (n_out, 1),
    },
    "debug_print_layer_output_template": True,
  })
  print("Creating network...")
  with tf.Graph().as_default():
    network = TFNetwork(config=config, train_flag=True)

    net_dict = {}
    layer_n_out = 10
    layer_common_args = {"class": "linear", "activation": "relu", "n_out": layer_n_out, "L2": 0.01}

    def layer(sources, **kwargs):
      args = kwargs.copy()
      for k, v in layer_common_args.items():
        args.setdefault(k, v)
      args.setdefault("from", sources)
      return args

    def make_network(num_layers):
      net_dict["input"] = {"class": "tikhonov_regularization", "meta_loss_scale": 0.1, "from": "data"}
      sources = ["input"]
      for i in range(num_layers):
        net_dict["layer%i" % i] = layer(sources=sources)
        sources = ["layer%i" % i]
      net_dict["output"] = {"class": "softmax", "loss": "ce", "loss_opts": {"use_fused": False}, "from": sources}

    make_network(num_layers=3)
    network.construct_from_dict(net_dict)
    data_input = network.extern_data.get_default_input_data()
    data_target = network.extern_data.get_default_target_data()
    from returnn.tf.updater import Updater
    updater = Updater(config=config, network=network, initial_learning_rate=0.001)
    updater.set_trainable_vars(tf_compat.v1.trainable_variables())
    update_op = updater.get_optim_op()
    assert updater.optim_meta_losses_dict
    fetches = network.get_fetches_dict()
    fetches.update(updater.optim_meta_losses_dict)

    n_batch = 5
    n_time = 11
    rnd = numpy.random.RandomState(42)
    with tf_compat.v1.Session() as session:
      session.run(tf_compat.v1.variables_initializer(tf_compat.v1.global_variables() + [network.global_train_step_var]))
      for step in range(5):
        info, _ = session.run(
          (fetches, update_op),
          feed_dict={
            data_input.placeholder: rnd.normal(size=(n_batch, n_time, n_in)).astype("float32"),
            data_input.size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
            data_target.placeholder: rnd.randint(0, n_out, size=(n_batch, n_time), dtype="int32"),
            data_target.size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
          })
        print("step:", step, "info:", info)


def test_split_info_input():
  from returnn.tf.util.basic import print_graph_output, find_ops_with_tensor_input
  config = Config({
    "debug_print_layer_output_template": True,
    "extern_data": {"data": {"dim": 7}}
  })
  net_dict = {
    "a": {"class": "linear", "activation": "tanh", "n_out": 11, "from": "data:data"},
    "b": {"class": "linear", "activation": "tanh", "n_out": 13, "from": "data:data"},
    "concat": {"class": "copy", "from": ["a", "b"]},
    "output": {"class": "linear", "activation": None, "with_bias": True, "from": ["concat"], "n_out": 17}
  }
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    out_weights = network.get_default_output_layer().params["W"]
    print("out_weights:", out_weights)
    assert isinstance(out_weights, tf.Variable)
    assert out_weights.get_shape().as_list() == [11 + 13, 17]
    # TODO: multiple checks:
    # the split info itself
    # the param init handling...
    # actually, for param init handling, input dim splits do not matter. they matter just for copying/growing-pretrain.
    # for param init handling, output dim split do matter.


def test_VariableLayer_split_info():
  feat1 = FeatureDim("feature1", 3)
  feat2 = FeatureDim("feature2", 5)
  net_dict = {
    "output": {"class": "variable", "shape": [2 * feat1 + feat2, 3 * feat1]}
  }
  with make_scope() as session:
    network = TFNetwork(extern_data=ExternData())
    network.construct_from_dict(net_dict)
    layer = network.layers["output"]
    assert isinstance(layer, VariableLayer)
    assert_equal(
      tf_util.get_param_axes_split_info(layer.output.placeholder),
      [2 * [feat1.dimension] + [feat2.dimension], 3 * [feat1.dimension]])


def test_VariableLayer_init_by_layer():
  # https://github.com/rwth-i6/returnn/wiki/Parameter-initialization
  # https://pytorch.org/docs/1.9.1/_modules/torch/nn/modules/linear.html#Linear
  # https://pytorch.org/docs/1.9.1/nn.init.html
  # https://github.com/pytorch/pytorch/blob/d665097cad3207795a655bbdde7a4123c0adc1c3/torch/nn/modules/linear.py#L96
  # https://github.com/pytorch/pytorch/issues/57109
  # https://github.com/pytorch/pytorch/blob/77721ee318d6785010144aa4569efb98199e7162/torch/nn/init.py#L390-L395
  # https://github.com/pytorch/pytorch/pull/41638
  # https://github.com/pytorch/pytorch/issues/18182
  """
  PyTorch new / proposed Linear weight:
    kaiming_normal_(mode='fan_out')  # default nonlinearity='leaky_relu', default neg slope 0.01
    (ignoring neg slope) -> std = sqrt(2 / fan_out)

  PyTorch old / current Linear weight:
    kaiming_uniform_(self.weight, a=math.sqrt(5))  # default nonlinearity='leaky_relu', mode='fan_in'
    -> std = sqrt(1 / (3 * fan_in))  -> bound = sqrt(1/fan_in)

  For uniform distribution, bound = sqrt(3) * std.
  For truncated normal distribution: stddev /= .87962566103423978  # TF VarianceScaling, scipy a=-2, b=2 ...
  """
  shape = (3, 4)
  net_dict = {
    "random": {"class": "random", "shape": shape, "distribution": "truncated_normal", "static": True},
    "var": {"class": "variable", "shape": shape, "init_by_layer": "random"},
    "output": {"class": "copy", "from": "var"}
  }
  config = Config({
    "extern_data": {"data": {"dim": 4}},  # not actually used...
  })
  tf_rnd_seed = 42
  with make_scope() as session:
    tf_compat.v1.set_random_seed(tf_rnd_seed)
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    assert_equal(net.layers["random"].params, {})
    assert_equal(net.get_params_list(), [next(iter(net.layers["var"].params.values()))])
    net.initialize_params(session)
    var_v = session.run(net.layers["var"].output.placeholder)
  # Run again to check that it is deterministic.
  with make_scope() as session:
    tf_compat.v1.set_random_seed(tf_rnd_seed)
    net = TFNetwork(config=config)
    net.construct_from_dict(net_dict)
    net.initialize_params(session)
    var_v_ = session.run(net.layers["var"].output.placeholder)
  numpy.testing.assert_array_equal(var_v, var_v_)


def test_extra1():
  n_in, n_out = 2, 3
  config = Config({
    "extern_data": {"data": {"dim": n_in}},
    "debug_print_layer_output_template": True,
  })
  net_dict = {
    "input": {"class": "linear", "activation": "relu", "n_out": n_out, "from": "data"},
    "extra.2:input": {"class": "linear", "activation": None, "n_out": n_out, "from": "data"},
    # "extra.3:input automatically ...
    "output1": {"class": "copy", "from": "input", "is_output_layer": True},
    "output2": {"class": "activation", "from": "extra.2:input", "activation": "relu", "is_output_layer": True},
    "output3": {"class": "copy", "from": "extra.3:input", "is_output_layer": True},
  }
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)

    assert "extra.2" in network.extra_nets
    assert "extra.3" in network.extra_nets
    params = network.get_params_list()
    print("Params:", params)
    assert len(params) == 2  # W + b

    feed_dict = make_feed_dict(network.extern_data)
    session.run(tf_compat.v1.variables_initializer(tf_compat.v1.global_variables() + [network.global_train_step_var]))
    out1 = session.run(network.layers["output1"].output.placeholder, feed_dict=feed_dict)
    out2 = session.run(network.layers["output2"].output.placeholder, feed_dict=feed_dict)
    out3 = session.run(network.layers["output3"].output.placeholder, feed_dict=feed_dict)
    numpy.testing.assert_almost_equal(out1, out2)
    numpy.testing.assert_almost_equal(out2, out3)


def test_extra_subnet():
  n_in, n_out = 3, 3
  config = Config({
    "extern_data": {"data": {"dim": n_in}},
    "debug_print_layer_output_template": True,
  })
  net_dict = {
    "subnet": {
      "class": "subnetwork",
      "from": "data:data",
      "subnetwork": {
        "output": {"class": "linear", "activation": "relu", "n_out": n_out, "from": "data"},
        "output2": {"class": "linear", "activation": "relu", "n_out": n_out, "from": "data", "is_output_layer": True},
      },
    },
    "extra.2:subnet": {
      "class": "subnetwork",
      "from": "data:data",
      "subnetwork": {
        "output": {"class": "copy", "from": "data"},
        "output2": {
          "class": "linear", "from": "data:data", "activation": None, "n_out": n_out, "is_output_layer": True},
      },
    },
    # extra.3:subnet automatically
    "sub1_output1": {"class": "copy", "from": "subnet/output", "is_output_layer": True},
    "sub1_output2": {"class": "copy", "from": "subnet/output2", "is_output_layer": True},
    "sub2_output1": {"class": "copy", "from": "extra.2:subnet/output", "is_output_layer": True},
    "sub2_output2": {
      "class": "activation", "activation": "relu", "from": "extra.2:subnet/output2", "is_output_layer": True},
    "sub3_output1": {"class": "copy", "from": "extra.3:subnet/output", "is_output_layer": True},
    "sub3_output2": {"class": "copy", "from": "extra.3:subnet/output2", "is_output_layer": True},
  }
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)
    assert "extra.2" in network.extra_nets
    assert "extra.3" in network.extra_nets
    params = network.get_params_list()
    print("Params:", params)
    assert len(params) == 4

    feed_dict = make_feed_dict(network.extern_data)
    session.run(tf_compat.v1.variables_initializer(tf_compat.v1.global_variables() + [network.global_train_step_var]))
    in_ = feed_dict[network.extern_data.data["data"].placeholder]
    sub1_out1 = session.run(network.layers["sub1_output1"].output.placeholder, feed_dict=feed_dict)
    sub1_out2 = session.run(network.layers["sub1_output2"].output.placeholder, feed_dict=feed_dict)
    sub2_out1 = session.run(network.layers["sub2_output1"].output.placeholder, feed_dict=feed_dict)
    sub2_out2 = session.run(network.layers["sub2_output2"].output.placeholder, feed_dict=feed_dict)
    sub3_out1 = session.run(network.layers["sub3_output1"].output.placeholder, feed_dict=feed_dict)
    sub3_out2 = session.run(network.layers["sub3_output2"].output.placeholder, feed_dict=feed_dict)
    numpy.testing.assert_almost_equal(sub1_out1, sub3_out1)
    numpy.testing.assert_almost_equal(sub1_out2, sub3_out2)
    numpy.testing.assert_almost_equal(sub1_out2, sub2_out2)
    numpy.testing.assert_almost_equal(in_, sub2_out1)


def test_subnetwork_unused_output():
  # https://github.com/rwth-i6/returnn/issues/580
  with make_scope() as session:
    net_dict = {
      'sub': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'linear': {'class': 'linear', 'from': 'base:data:data', 'n_out': 1},
          'linear_0': {'class': 'linear', 'from': 'linear', 'n_out': 1},
          'output': {'class': 'copy', 'from': 'linear'}}},
      'linear': {'class': 'linear', 'from': ['sub/linear', 'sub/linear_0'], 'n_out': 1},
      'output': {'class': 'copy', 'from': 'linear'}}
    config = Config()
    config.update(dict(num_inputs=1, num_outputs=1))
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)


def test_subnetwork_deep_stack():
  # https://github.com/rwth-i6/returnn/issues/993
  if not util.PY3:
    raise unittest.SkipTest("test case needs python 3")  # __qualname__, __func__, etc
  with make_scope() as session:
    import better_exchook
    import traceback
    import types
    import collections

    max_depth = 5

    def _find_net_construct_from_dict_in_stack(stack):
      found_last = None
      for i, frame in enumerate(stack):
        assert isinstance(frame, types.FrameType)
        func = better_exchook.get_func_from_code_object(frame.f_code)
        if func == TFNetwork.construct_from_dict:
          found_last = i
      if found_last is not None:
        return found_last
      raise Exception("construct_from_dict not found in stack")

    def _extract_stack():
      stack = [frame for frame, _ in traceback.walk_stack(None)]
      stack = stack[:_find_net_construct_from_dict_in_stack(stack) + 1]
      return stack

    def _top_eval(source, **_):
      print("Top eval func stack:")
      stack = _extract_stack()
      # LayerBase.transform_config_dict would only increase if this is not flat net construction,
      # see https://github.com/rwth-i6/returnn/issues/992.
      # For reference, output without fixes: https://gist.github.com/albertz/a1a20710e4d08292e35ba6910ca132e8
      func_counter = collections.Counter()
      for frame in stack:
        assert isinstance(frame, types.FrameType)
        func = better_exchook.get_func_from_code_object(frame.f_code)
        if not func:
          print("  Warning: unexpected code object: %s" % frame.f_code.co_name)
          continue
        func_counter[func] += 1
        print(" ", func.__qualname__)
      print("Stack depth:", len(stack))
      print("Num functions:", len(func_counter))
      for func, count in func_counter.most_common(10):
        print("Most common func %s: count %i" % (func.__qualname__, count))
      """
      For reference, without flat construction (#992), before the fix of #993, I get:
        Most common func Subnetwork.get_sub_layer_func.<locals>.wrapped_get_layer: count 30
        Most common func Subnetwork.get_layer_func.<locals>.wrapped_get_layer: count 30
        Most common func TFNetwork.construct_layer: count 12
        Most common func TFNetwork.construct_layer.<locals>.get_layer: count 11
        Most common func TFNetwork.add_layer: count 6
        Most common func LayerBase.transform_config_dict.<locals>.<listcomp>: count 6
        Most common func LayerBase.transform_config_dict: count 6
        Most common func CopyLayer.transform_config_dict: count 6
        Most common func Subnetwork.construct_layer: count 5
        Most common func SubnetworkLayer.transform_config_dict: count 5
      """
      print("Num TFNetwork.construct_layer:", func_counter[TFNetwork.construct_layer])
      print("Num LayerBase.transform_config_dict:", func_counter[LayerBase.transform_config_dict.__func__])
      # Before the fix, the stack depth was 125; with the improved variant (even not flat) it is 65.
      # We don't check for the exact number to allow for some variation of future changes.
      # However, we want to avoid that it becomes too deep in any case.
      assert len(stack) <= 70
      return source(0)

    def _create_subnet_layer_dict(depth):
      if depth >= max_depth:
        return {"class": "eval", "from": "base:" * depth + "data", "eval": _top_eval}
      return {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'sub%i' % depth: _create_subnet_layer_dict(depth + 1),
          'output': {'class': 'copy', 'from': 'sub%i' % depth}}}

    net_dict = {
      'sub': _create_subnet_layer_dict(0),
      'output': {'class': 'copy', 'from': 'sub'}}
    config = Config({"extern_data": {"data": {"dim": 3}}})
    network = TFNetwork(config=config)
    network.construct_from_dict(net_dict)


def test_subnet_construct_layer():
  # https://github.com/rwth-i6/returnn/issues/1014
  with make_scope() as session:
    net = TFNetwork(config=Config({"extern_data": {"data": {"dim": 1}}}))
    subnet_layer = net.construct_layer(name="subnet", net_dict={
      "subnet": {"class": "subnetwork", "from": "data", "subnetwork": {"output": {"class": "copy", "from": "data"}}}})
    assert isinstance(subnet_layer, SubnetworkLayer)
    assert subnet_layer.output.placeholder is net.extern_data.get_default_input_data().placeholder


def test_extra_search():
  class Callbacks:
    history = []
    @classmethod
    def callback(cls, self, source, **kwargs):
      """
      :param LayerBase self:
      :param (int)->tf.Tensor source:
      :rtype: tf.Tensor
      """
      print("test_extra_search, callback: %r, %r; search flag %r" % (
        self.network.name, self, self.network.search_flag))
      cls.history.append(self)
      return source(0)

  n_batch, n_time, n_in, n_out = 2, 3, 7, 11
  rnd = numpy.random.RandomState(42)
  config = Config({
    "debug_print_layer_output_template": True,
    "extern_data": {"data": {"dim": n_in}}
  })
  net_dict = {
    "input": {"class": "eval", "eval": Callbacks.callback, "from": "search_post_output"},
    "extra.search:input": {"class": "eval", "eval": Callbacks.callback, "from": "data"},
    # Note: This 'output' layer is created twice: Once in main net, once in extra-net.
    "output": {"class": "subnetwork", "from": "input", "subnetwork": {
      "inner": {"class": "linear", "from": "data", "activation": "relu", "n_out": n_out},
      "output": {"class": "eval", "from": "inner", "eval": Callbacks.callback}
    }},
    "search_post_output": {"class": "linear", "from": "extra.search:output", "activation": "relu", "n_out": n_in}
  }
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(net_dict)

    assert not network.search_flag
    assert "extra.search" in network.extra_nets
    extra_net = network.extra_nets["extra.search"]
    assert extra_net.search_flag
    assert "input" in network.layers
    assert "input" in extra_net.layers
    assert "output" in network.layers
    assert "output" in extra_net.layers
    layer_input = network.layers["input"]
    assert isinstance(layer_input, EvalLayer)
    assert layer_input in Callbacks.history
    assert layer_input.network is network
    layer_extra_input = extra_net.layers["input"]
    assert isinstance(layer_extra_input, EvalLayer)
    assert layer_extra_input in Callbacks.history
    assert layer_extra_input.network is extra_net
    layer_output = network.layers["output"]
    assert isinstance(layer_output, SubnetworkLayer)
    assert layer_output.network is network
    layer_output_output = layer_output.subnetwork.layers["output"]
    assert layer_output_output in Callbacks.history
    layer_extra_output = extra_net.layers["output"]
    assert isinstance(layer_extra_output, SubnetworkLayer)
    assert layer_extra_output.network is extra_net
    layer_extra_output_output = layer_extra_output.subnetwork.layers["output"]
    assert layer_extra_output_output in Callbacks.history

    fetches = network.get_fetches_dict()
    data_input = network.extern_data.data["data"]

    session.run(tf_compat.v1.variables_initializer(tf_compat.v1.global_variables() + [network.global_train_step_var]))
    info, out = session.run(
      (fetches, layer_output.output.placeholder),
      feed_dict={
        data_input.placeholder: rnd.normal(size=(n_batch, n_time, n_in)).astype("float32"),
        data_input.size_placeholder[0]: numpy.array([n_time] * n_batch, dtype="int32"),
      })
    print(info)
    print(out)  # random...


def test_HDFDumpLayer():
  import os
  from test_HDFDataset import get_test_tmp_file, DatasetTestReader, HDFDataset
  hdf_filename = get_test_tmp_file(".hdf")
  os.remove(hdf_filename)  # HDFDumpLayer expects that the file does not exist

  with make_scope() as session:
    n_in, n_out = 4, 3
    config = Config()
    config.update({
      "num_outputs": n_out,
      "num_inputs": n_in,
      "network": {
        "lstm": {"class": "rec", "unit": "LSTMBlock", "from": ["data"], "n_out": n_out},
        "dump": {"class": "hdf_dump", "filename": hdf_filename, "from": ["lstm"]},
        "output": {"class": "copy", "from": ["dump"]},
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))

    session.run(tf_compat.v1.global_variables_initializer())
    out = network.layers["output"].output.placeholder
    n_batch = 1
    seq_len = 4
    input_data = numpy.array([[
      [1, -0.2, 0.3, -4],
      [2, -0.6, 0.7, -1.8],
      [1, 0.3, -0.1, -0.8],
      [0.1, -0.2, 0.2, .8]]],
      dtype="float32")
    input_tags = numpy.array([b"seq-0"], dtype="S5")
    seq_lens = numpy.array([seq_len], dtype="int32")
    assert input_data.shape == (n_batch, seq_lens[0], n_in)
    feed = {network.extern_data.data["data"].placeholder: input_data,
            network.extern_data.data["data"].size_placeholder[0]: seq_lens,
            network.extern_data.data["seq_tag"].placeholder: input_tags}
    assert_equal(feed[network.extern_data.get_default_input_data().placeholder].shape, (n_batch, seq_len, n_in))
    session.run([out, network.get_post_control_dependencies()], feed_dict=feed)

    network.call_graph_reset_callbacks()

  assert os.path.exists(hdf_filename)
  reader = DatasetTestReader(HDFDataset([hdf_filename]))
  reader.read_all()
  assert reader.num_seqs == 1
  assert reader.seq_tags == ["seq-0"]
  assert_equal(reader.seq_lens[0]["data"], seq_lens[0])
  assert_equal(reader.data["data"][0].shape, (seq_lens[0], n_out))


def test_HDFDumpLayer_sparse():
  import os
  from test_HDFDataset import get_test_tmp_file, DatasetTestReader, HDFDataset
  hdf_filename = get_test_tmp_file(".hdf")
  os.remove(hdf_filename)  # HDFDumpLayer expects that the file does not exist

  with make_scope() as session:
    n_in, n_out = 4, 5
    config = Config()
    config.update({
      "num_inputs": n_in,
      "num_outputs": n_out,
      "network": {
        "dump": {
          "class": "hdf_dump", "filename": hdf_filename, "from": "data:classes",
          "is_output_layer": True
        },
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))

    session.run(tf_compat.v1.global_variables_initializer())
    n_batch = 1
    classes_data = numpy.array([[2, 5, 6]], dtype="int32")
    classes_seq_lens = [classes_data.shape[1]]
    assert classes_data.shape == (n_batch, classes_seq_lens[0])
    input_tags = numpy.array([b"seq-0"], dtype="S5")
    feed = {network.extern_data.data["classes"].placeholder: classes_data,
            network.extern_data.data["classes"].size_placeholder[0]: classes_seq_lens,
            network.extern_data.data["seq_tag"].placeholder: input_tags}
    session.run(network.get_fetches_dict(), feed_dict=feed)

    network.call_graph_reset_callbacks()

  assert os.path.exists(hdf_filename)
  reader = DatasetTestReader(HDFDataset([hdf_filename]))
  reader.read_all()
  assert reader.num_seqs == 1
  assert reader.seq_tags == ["seq-0"]
  assert_equal(reader.seq_lens[0]["data"], classes_seq_lens[0])
  assert_equal(reader.data["data"][0].shape, (classes_seq_lens[0],))
  assert_equal(reader.data_sparse["data"], True)
  assert_equal(reader.dataset.get_data_dim("data"), n_out)


def test_HDFDumpLayer_fixed_length():
  import os
  from test_HDFDataset import get_test_tmp_file, DatasetTestReader, HDFDataset
  hdf_filename = get_test_tmp_file(".hdf")
  os.remove(hdf_filename)  # HDFDumpLayer expects that the file does not exist

  with make_scope() as session:
    n_in, n_out = 4, 3
    config = Config()
    config.update({
      "num_outputs": n_out,
      "num_inputs": n_in,
      "network": {
        "lstm": {"class": "rec", "unit": "LSTMBlock", "from": ["data"], "n_out": n_out},
        "last_state": {"class": "get_last_hidden_state", "from": ["lstm"], "key": "h", "n_out": n_out},
        "last_state_expanded": {"class": "expand_dims", "from": ["last_state"], "axis": "T"},
        "dump": {"class": "hdf_dump", "filename": hdf_filename, "from": ["last_state_expanded"]},
        "output": {"class": "copy", "from": ["dump"]},
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))

    session.run(tf_compat.v1.global_variables_initializer())
    out = network.layers["output"].output.placeholder
    n_batch = 1
    seq_len = 4
    input_data = numpy.array([[
      [1, -0.2, 0.3, -4],
      [2, -0.6, 0.7, -1.8],
      [1, 0.3, -0.1, -0.8],
      [0.1, -0.2, 0.2, .8]]],
      dtype="float32")
    input_tags = numpy.array([b"seq-0"], dtype="S5")
    seq_lens = numpy.array([seq_len], dtype="int32")
    assert input_data.shape == (n_batch, seq_lens[0], n_in)
    feed = {network.extern_data.data["data"].placeholder: input_data,
            network.extern_data.data["data"].size_placeholder[0]: seq_lens,
            network.extern_data.data["seq_tag"].placeholder: input_tags}
    session.run([out, network.get_post_control_dependencies()], feed_dict=feed)

    network.call_graph_reset_callbacks()

  assert os.path.exists(hdf_filename)
  reader = DatasetTestReader(HDFDataset([hdf_filename]))
  reader.read_all()
  assert reader.num_seqs == 1
  assert reader.seq_tags == ["seq-0"]
  assert_equal(reader.seq_lens[0]["data"], 1)
  assert_equal(reader.data["data"][0].shape, (1, n_out))


def test_HDFDumpLayer_extra():
  import os
  from test_HDFDataset import get_test_tmp_file, DatasetTestReader, HDFDataset
  hdf_filename = get_test_tmp_file(".hdf")
  os.remove(hdf_filename)  # HDFDumpLayer expects that the file does not exist

  with make_scope() as session:
    n_in = 5
    n_out1 = 7
    config = Config()
    config.update({
      "extern_data": {
        "data": {"dim": n_in},
        "classes1": {"dim": n_out1, "sparse": True},
        "classes2": {"dim": None, "dtype": "float32", "shape": ()},
      },
      "network": {
        "dump": {
          "class": "hdf_dump", "filename": hdf_filename,
          "from": "data",
          "extra": {"classes1": "data:classes1", "classes2": "data:classes2"},
          "is_output_layer": True
        },
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    network.print_network_info()

    session.run(tf_compat.v1.global_variables_initializer())
    n_batch = 1
    input_data = numpy.array([[
      [1, -0.2, 0.3, -4, 5],
      [2, -0.6, 0.7, -1.8, 2.9],
      [1, 0.3, -0.1, -0.8, 0.5],
      [0.1, -0.2, 0.2, .8, -0.3]]],
      dtype="float32")
    input_seq_lens = [input_data.shape[1]]
    assert input_data.shape == (n_batch, input_seq_lens[0], n_in)
    classes1_data = numpy.array([[2, 5, 6]], dtype="int32")
    classes1_seq_lens = [classes1_data.shape[1]]
    assert classes1_data.shape == (n_batch, classes1_seq_lens[0])
    classes2_data = numpy.array([-7.89], dtype="float32")
    assert classes2_data.shape == (n_batch,)
    seq_tags = numpy.array([b"seq-0"], dtype="S5")
    feed = {
      network.extern_data.data["data"].placeholder: input_data,
      network.extern_data.data["data"].size_placeholder[0]: input_seq_lens,
      network.extern_data.data["classes1"].placeholder: classes1_data,
      network.extern_data.data["classes1"].size_placeholder[0]: classes1_seq_lens,
      network.extern_data.data["classes2"].placeholder: classes2_data,
      network.extern_data.data["seq_tag"].placeholder: seq_tags}
    fetches = network.get_fetches_dict()
    result = session.run(fetches, feed_dict=feed)
    pprint(result)

    network.call_graph_reset_callbacks()

  assert os.path.exists(hdf_filename)
  reader = DatasetTestReader(HDFDataset([hdf_filename]))
  reader.read_all()
  assert reader.num_seqs == 1
  assert reader.seq_tags == ["seq-0"]
  assert_equal(reader.seq_lens[0]["data"], input_seq_lens[0])
  assert_equal(reader.data["data"][0].shape, (input_seq_lens[0], n_in))
  assert_equal(reader.data["classes1"][0].shape, (classes1_seq_lens[0],))
  assert_equal(reader.data["classes2"][0].shape, (1,))
  numpy.testing.assert_almost_equal(reader.data["data"][0], input_data[0])
  numpy.testing.assert_equal(reader.data["classes1"][0], classes1_data[0])
  numpy.testing.assert_equal(reader.data["classes2"][0], [classes2_data[0]])


def test_HDFDumpLayer_dump_whole_batch_extra_sm():
  import os
  from test_HDFDataset import get_test_tmp_file, DatasetTestReader, HDFDataset
  hdf_filename = get_test_tmp_file(".hdf")
  os.remove(hdf_filename)  # HDFDumpLayer expects that the file does not exist
  rnd = numpy.random.RandomState(42)

  with make_scope() as session:
    n_in = 5
    config = Config()
    config.update({
      "extern_data": {
        "data": {"dim": n_in},
        "sm": dict(shape=(None, None)),
      },
      "network": {
        "dump": {
          "class": "hdf_dump", "filename": hdf_filename,
          "from": "data",
          "extra": {"sm": "data:sm"},
          "is_output_layer": True,
          "dump_whole_batches": True,
        },
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    network.print_network_info()

    session.run(tf_compat.v1.global_variables_initializer())
    n_batch = 1
    input_data = numpy.array([[
      [1, -0.2, 0.3, -4, 5],
      [2, -0.6, 0.7, -1.8, 2.9],
      [1, 0.3, -0.1, -0.8, 0.5],
      [0.1, -0.2, 0.2, .8, -0.3]]],
      dtype="float32")
    input_seq_lens = [input_data.shape[1]]
    assert input_data.shape == (n_batch, input_seq_lens[0], n_in)
    sm_seq_lens1 = [13]
    sm_seq_lens2 = [17]
    sm_data = rnd.normal(size=(n_batch, sm_seq_lens1[0], sm_seq_lens2[0])).astype(dtype="float32")
    seq_tags = numpy.array([b"seq-0"], dtype="S5")
    feed = {
      network.extern_data.data["data"].placeholder: input_data,
      network.extern_data.data["data"].size_placeholder[0]: input_seq_lens,
      network.extern_data.data["sm"].placeholder: sm_data,
      network.extern_data.data["sm"].size_placeholder[0]: sm_seq_lens1,
      network.extern_data.data["sm"].size_placeholder[1]: sm_seq_lens2,
      network.extern_data.data["seq_tag"].placeholder: seq_tags}
    fetches = network.get_fetches_dict()
    result = session.run(fetches, feed_dict=feed)
    pprint(result)

    network.call_graph_reset_callbacks()

  assert os.path.exists(hdf_filename)
  reader = DatasetTestReader(HDFDataset([hdf_filename]))
  reader.read_all()
  assert reader.num_seqs == 1
  assert reader.seq_tags == ["seq-0"]
  assert_equal(reader.seq_lens[0]["data"], input_seq_lens[0])
  assert_equal(reader.data["data"][0].shape, (input_seq_lens[0], n_in))
  numpy.testing.assert_almost_equal(reader.data["data"][0], input_data[0])
  assert_equal(reader.data["sm"][0].shape, (sm_seq_lens1[0] * sm_seq_lens2[0],))
  numpy.testing.assert_equal(numpy.reshape(reader.data["sm"][0], sm_data[0].shape), sm_data[0])


def test_HDFDumpLayer_dump_whole_batch_extra_sm1():
  import os
  from test_HDFDataset import get_test_tmp_file, DatasetTestReader, HDFDataset
  hdf_filename = get_test_tmp_file(".hdf")
  os.remove(hdf_filename)  # HDFDumpLayer expects that the file does not exist
  rnd = numpy.random.RandomState(42)

  with make_scope() as session:
    n_in = 5
    config = Config()
    config.update({
      "extern_data": {
        "data": {"dim": n_in},
        "sm": dict(shape=(None, 1, None), batch_dim_axis=1, feature_dim_axis=2),
      },
      "network": {
        "dump": {
          "class": "hdf_dump", "filename": hdf_filename,
          "from": "data",
          "extra": {"sm": "data:sm"},
          "is_output_layer": True,
          "dump_whole_batches": True,
        },
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    network.print_network_info()

    session.run(tf_compat.v1.global_variables_initializer())
    n_batch = 1
    input_data = numpy.array([[
      [1, -0.2, 0.3, -4, 5],
      [2, -0.6, 0.7, -1.8, 2.9],
      [1, 0.3, -0.1, -0.8, 0.5],
      [0.1, -0.2, 0.2, .8, -0.3]]],
      dtype="float32")
    input_seq_lens = [input_data.shape[1]]
    assert input_data.shape == (n_batch, input_seq_lens[0], n_in)
    sm_seq_lens1 = [13]
    sm_seq_lens2 = [17]
    sm_data = rnd.normal(size=(sm_seq_lens1[0], n_batch, 1, sm_seq_lens2[0])).astype(dtype="float32")
    seq_tags = numpy.array([b"seq-0"], dtype="S5")
    feed = {
      network.extern_data.data["data"].placeholder: input_data,
      network.extern_data.data["data"].size_placeholder[0]: input_seq_lens,
      network.extern_data.data["sm"].placeholder: sm_data,
      network.extern_data.data["sm"].size_placeholder[0]: sm_seq_lens1,
      network.extern_data.data["sm"].size_placeholder[2]: sm_seq_lens2,
      network.extern_data.data["seq_tag"].placeholder: seq_tags}
    fetches = network.get_fetches_dict()
    result = session.run(fetches, feed_dict=feed)
    pprint(result)

    network.call_graph_reset_callbacks()

  assert os.path.exists(hdf_filename)
  reader = DatasetTestReader(HDFDataset([hdf_filename]))
  reader.read_all()
  assert reader.num_seqs == 1
  assert reader.seq_tags == ["seq-0"]
  assert_equal(reader.data["data"][0].shape, (input_seq_lens[0], n_in))
  numpy.testing.assert_almost_equal(reader.data["data"][0], input_data[0])
  assert_equal(reader.data["sm"][0].shape, (sm_seq_lens1[0] * sm_seq_lens2[0],))
  sm_data_ = numpy.transpose(sm_data, (1, 0, 3, 2))
  numpy.testing.assert_equal(numpy.reshape(reader.data["sm"][0], sm_data_[0].shape), sm_data_[0])


def test_CrossEntropyLoss():
  with make_scope() as session:
    n_out = 13
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
        "classes": {"dim": n_out, "sparse": True},
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "var": {"class": "variable", "shape": (n_out,)},
      "add": {"class": "combine", "kind": "add", "from": ["data", "var"]},
      "output": {
        "class": "activation", "from": "add", "activation": "softmax",
        "loss": "ce"},
    })
    losses_dict, total_loss, total_constraints = net.get_losses_initialized()
    print("Losses:")
    pprint(losses_dict)
    assert set(losses_dict.keys()) == {"output"}
    loss_holder = losses_dict["output"]
    assert isinstance(loss_holder, LossHolder)
    assert isinstance(loss_holder.loss, CrossEntropyLoss)
    session.run(tf_compat.v1.global_variables_initializer())
    print("Get loss:")
    feed_dict = make_feed_dict(net.extern_data.data.values(), same_time=True)
    print("random classes:", feed_dict[net.extern_data.data["classes"].placeholder])
    loss_t = loss_holder.get_loss_value()
    opt = tf_compat.v1.train.GradientDescentOptimizer(learning_rate=0.1)
    minimize_op = opt.minimize(loss_t)
    last_loss_v = float("inf")
    for step in range(3):
      loss_v, _ = session.run((loss_t, minimize_op), feed_dict=feed_dict)
      print("step %i, loss %f" % (step, loss_v))
      assert numpy.isfinite(loss_v) and numpy.isscalar(loss_v)
      assert loss_v < last_loss_v  # it's convex and we cannot overshoot
      last_loss_v = loss_v


def test_CrossEntropyLoss_masked_inf():
  with make_scope() as session:
    n_out = 13
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
        "classes": {"dim": n_out, "sparse": True},
      }})
    mask_t = tf_compat.v1.placeholder(tf.bool, (n_out,), name="mask")

    def mask_func(source, **kwargs):
      x = source(0)
      assert x.shape.ndims == 3  # (B,T,n_out)
      from returnn.tf.util.basic import where_bc
      mask_bc = mask_t[None, None, :]  # (1,1,n_out)
      return where_bc(mask_bc, x, float("-inf"))

    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "var": {"class": "variable", "shape": (n_out,)},  # such that we can check that there are no nan/inf grads
      "add": {"class": "combine", "kind": "add", "from": ["data", "var"]},
      "mask": {"class": "eval", "from": "add", "eval": mask_func},
      "output": {
        "class": "activation", "from": "mask", "activation": "softmax",
        "loss": "ce"},
    })
    losses_dict, total_loss, total_constraints = net.get_losses_initialized()
    print("Losses:")
    pprint(losses_dict)
    assert set(losses_dict.keys()) == {"output"}
    loss_holder = losses_dict["output"]
    assert isinstance(loss_holder, LossHolder)
    assert isinstance(loss_holder.loss, CrossEntropyLoss)
    session.run(tf_compat.v1.global_variables_initializer())
    print("Get loss:")
    feed_dict = make_feed_dict(net.extern_data.data.values(), same_time=True)
    mask_v = numpy.array([True] * n_out)
    feed_dict[mask_t] = mask_v
    loss_t = loss_holder.get_loss_value()
    opt = tf_compat.v1.train.GradientDescentOptimizer(learning_rate=0.1)
    minimize_op = opt.minimize(loss_t)
    last_loss_v = float("inf")
    for step in range(3):
      loss_v, _ = session.run((loss_t, minimize_op), feed_dict=feed_dict)
      print("step %i, loss %f" % (step, loss_v))
      assert numpy.isfinite(loss_v) and numpy.isscalar(loss_v)
      assert loss_v < last_loss_v  # it's convex and we cannot overshoot
      last_loss_v = loss_v
    print("Now mask.")
    feed_dict = make_feed_dict(net.extern_data.data.values(), same_time=True, n_batch=1, n_time=1)
    feed_dict[mask_t] = mask_v
    rnd_classes = feed_dict[net.extern_data.data["classes"].placeholder]
    print("random classes:", rnd_classes)
    mask_v[rnd_classes[0, 0]] = False
    var_t, = tf_compat.v1.trainable_variables()
    last_var_v = session.run(var_t)
    for step in range(3, 6):
      loss_v, _ = session.run((loss_t, minimize_op), feed_dict=feed_dict)
      print("step %i, loss %f" % (step, loss_v))
      assert numpy.isinf(loss_v) and numpy.isscalar(loss_v)
      var_v = session.run(var_t)
      assert numpy.isfinite(var_v).all()  # while the loss is inf, the gradients should be finite!
      assert not (var_v == last_var_v).all()  # and there also was some non-zero gradient!
      last_var_v = var_v


def test_CrossEntropyLoss_masked_inf_fake_upper_bound():
  # Almost the same as test_CrossEntropyLoss_masked_inf, but we use fake_upper_bound.
  with make_scope() as session:
    n_out = 13
    fake_upper_bound = 10.
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
        "classes": {"dim": n_out, "sparse": True},
      }})
    mask_t = tf_compat.v1.placeholder(tf.bool, (n_out,), name="mask")

    def mask_func(source, **kwargs):
      x = source(0)
      assert x.shape.ndims == 3  # (B,T,n_out)
      from returnn.tf.util.basic import where_bc
      mask_bc = mask_t[None, None, :]  # (1,1,n_out)
      return where_bc(mask_bc, x, float("-inf"))

    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "var": {"class": "variable", "shape": (n_out,)},  # such that we can check that there are no nan/inf grads
      "add": {"class": "combine", "kind": "add", "from": ["data", "var"]},
      "mask": {"class": "eval", "from": "add", "eval": mask_func},
      "output": {
        "class": "activation", "from": "mask", "activation": "softmax",
        "loss": "ce", "loss_opts": {"fake_upper_bound": fake_upper_bound}},
    })
    losses_dict, total_loss, total_constraints = net.get_losses_initialized()
    print("Losses:")
    pprint(losses_dict)
    assert set(losses_dict.keys()) == {"output"}
    loss_holder = losses_dict["output"]
    assert isinstance(loss_holder, LossHolder)
    assert isinstance(loss_holder.loss, CrossEntropyLoss)
    session.run(tf_compat.v1.global_variables_initializer())
    print("Get loss:")
    feed_dict = make_feed_dict(net.extern_data.data.values(), same_time=True)
    mask_v = numpy.array([True] * n_out)
    feed_dict[mask_t] = mask_v
    loss_t = loss_holder.get_loss_value()
    opt = tf_compat.v1.train.GradientDescentOptimizer(learning_rate=0.1)
    minimize_op = opt.minimize(loss_t)
    last_loss_v = float("inf")
    for step in range(3):
      loss_v, _ = session.run((loss_t, minimize_op), feed_dict=feed_dict)
      print("step %i, loss %f" % (step, loss_v))
      assert numpy.isfinite(loss_v) and numpy.isscalar(loss_v)
      assert loss_v < last_loss_v  # it's convex and we cannot overshoot
      last_loss_v = loss_v
    print("Now mask.")
    feed_dict = make_feed_dict(net.extern_data.data.values(), same_time=True, n_batch=1, n_time=1)
    feed_dict[mask_t] = mask_v
    rnd_classes = feed_dict[net.extern_data.data["classes"].placeholder]
    print("random classes:", rnd_classes)
    mask_v[rnd_classes[0, 0]] = False
    var_t, = tf_compat.v1.trainable_variables()
    last_var_v = session.run(var_t)
    for step in range(3, 6):
      loss_v, _ = session.run((loss_t, minimize_op), feed_dict=feed_dict)
      print("step %i, loss %f" % (step, loss_v))
      assert loss_v == fake_upper_bound and numpy.isscalar(loss_v)
      var_v = session.run(var_t)
      assert numpy.isfinite(var_v).all()  # while the loss is bounded, the gradients should be finite!
      assert not (var_v == last_var_v).all()  # and there also was some non-zero gradient!
      last_var_v = var_v


def test_MeanSquaredError():
  with make_scope() as session:
    n_out = 13
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
        "classes": {"dim": n_out, "sparse": False, "shape": (1, n_out), "time_dim_axis": 1},
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "var": {"class": "variable", "shape": (n_out,)},
      "add": {"class": "combine", "kind": "add", "from": ["data", "var"]},
      "reduce": {"class": "reduce", "mode": "mean", "axes": "T", "from": "add", "keep_dims": True},
      "output": {
        "class": "activation", "from": "reduce", "activation": "sigmoid",
        "loss": "mse"},
    })
    losses_dict, total_loss, total_constraints = net.get_losses_initialized()
    print("Losses:")
    pprint(losses_dict)
    assert set(losses_dict.keys()) == {"output"}
    loss_holder = losses_dict["output"]
    assert isinstance(loss_holder, LossHolder)
    assert isinstance(loss_holder.loss, MeanSquaredError)
    session.run(tf_compat.v1.global_variables_initializer())
    print("Get loss:")
    feed_dict = make_feed_dict(net.extern_data.data.values(), same_time=True, n_time=1)
    print("random classes:", feed_dict[net.extern_data.data["classes"].placeholder])
    loss_t = loss_holder.get_loss_value()
    error_t = loss_holder.get_error_value()
    opt = tf_compat.v1.train.GradientDescentOptimizer(learning_rate=0.1)
    minimize_op = opt.minimize(loss_t)
    last_loss_v = float("inf")
    for step in range(3):
      loss_v, error_v, _ = session.run((loss_t, error_t, minimize_op), feed_dict=feed_dict)
      print("step %i, loss %f, error %f" % (step, loss_v, error_v))
      assert numpy.isfinite(loss_v) and numpy.isscalar(loss_v)
      assert loss_v <= last_loss_v
      last_loss_v = loss_v


def test_reduce_mean_in_time():
  with make_scope() as session:
    n_out = 5
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "output": {"class": "reduce", "mode": "mean", "axis": "T", "from": ["data"]}
    })
    session.run(tf_compat.v1.global_variables_initializer())
    out = net.layers["output"].output.placeholder
    n_batch = 3
    max_seq_len = 10
    feed = make_feed_dict(net.extern_data.data.values(), n_batch=n_batch, n_time=max_seq_len, same_time=True)
    v = session.run(out, feed_dict=feed)
    input_len = feed[net.extern_data.data["data"].size_placeholder[0]]
    input_data = feed[net.extern_data.data["data"].placeholder]

    ref = numpy.zeros([n_batch, n_out])
    for batch, seq_len in enumerate(input_len):
      ref[batch, :] = numpy.mean(input_data[batch, :seq_len, :], axis=0)

    numpy.testing.assert_allclose(ref, v, rtol=1e-5)


def test_reduce_mean_batch_time():
  with make_scope() as session:
    n_out = 5
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "output": {"class": "reduce", "mode": "mean", "axis": ["B", "T"], "from": ["data"]}
    })
    session.run(tf_compat.v1.global_variables_initializer())
    out = net.layers["output"].output.placeholder
    n_batch = 3
    max_seq_len = 10
    feed = make_feed_dict(net.extern_data.data.values(), n_batch=n_batch, n_time=max_seq_len, same_time=True)
    v = session.run(out, feed_dict=feed)
    input_len = feed[net.extern_data.data["data"].size_placeholder[0]]
    input_data = feed[net.extern_data.data["data"].placeholder]

    input_data_masked = numpy.copy(input_data)
    for batch, seq_len in enumerate(input_len):
      input_data_masked[batch, seq_len:, :] = numpy.nan
    ref = numpy.nanmean(input_data_masked, axis=(0, 1))

    numpy.testing.assert_allclose(ref, v, rtol=1e-5)


def test_automatic_seq_lengths():
  with make_scope() as session:
    n_out = 5
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "layer0": {
        'class': 'pad', 'mode': 'reflect', 'axes': 'spatial', 'padding': (3, 3), 'from': 'data'},  # len+6
      "layer1": {
        'class': 'conv', 'from': 'layer0', 'activation': None, 'with_bias': True, 'n_out': n_out,
        'filter_size': (7,), 'padding': 'valid', 'strides': (1,), 'dilation_rate': (1,)},  # max(len+6-6,0)
      "output": {"class": "copy", "from": "layer1"},
    })
    session.run(tf_compat.v1.global_variables_initializer())
    in_data = net.extern_data.get_default_input_data()
    out_data = net.layers["output"].output.copy_as_batch_spatial_major()
    assert_equal(out_data.shape, in_data.shape)
    n_batch = 3
    max_seq_len = 10
    feed = make_feed_dict([in_data], n_batch=n_batch, n_time=max_seq_len)
    out_lens = out_data.get_sequence_lengths()
    out_v, out_lens_v = session.run((out_data.placeholder, out_lens), feed_dict=feed)
    in_v = feed[in_data.placeholder]
    in_lens_v = feed[in_data.size_placeholder[0]]
    assert_equal(in_v.shape, out_v.shape)
    assert_equal(in_lens_v.tolist(), out_lens_v.tolist())
    # So far, everything should always be true, unless we have messed some op really up.
    # Now we want to do the main test, i.e. whether we get the same tensor.
    from returnn.tf.util.basic import print_graph_output
    print_graph_output(out_lens)
    assert out_lens is in_data.size_placeholder[0]


def test_automatic_seq_lengths2():
  with make_scope() as session:
    n_out = 5
    config = Config({
      "debug_print_layer_output_template": True,
      "extern_data": {
        "data": {"dim": n_out},
      }})
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict({
      "layer0": {
        'class': 'conv', 'from': 'data', 'activation': None, 'with_bias': True, 'n_out': n_out,
        'filter_size': (1,), 'padding': 'valid'},
      "output": {"class": "copy", "from": "layer0"},
    })
    session.run(tf_compat.v1.global_variables_initializer())
    in_data = net.extern_data.get_default_input_data()
    out_data = net.layers["output"].output.copy_as_batch_spatial_major()
    assert_equal(out_data.shape, in_data.shape)
    n_batch = 3
    max_seq_len = 10
    feed = make_feed_dict([in_data], n_batch=n_batch, n_time=max_seq_len)
    out_lens = out_data.get_sequence_lengths()
    out_v, out_lens_v = session.run((out_data.placeholder, out_lens), feed_dict=feed)
    in_v = feed[in_data.placeholder]
    in_lens_v = feed[in_data.size_placeholder[0]]
    assert_equal(in_v.shape, out_v.shape)
    assert_equal(in_lens_v.tolist(), out_lens_v.tolist())
    # So far, everything should always be true, unless we have messed some op really up.
    # Now we want to do the main test, i.e. whether we get the same tensor.
    from returnn.tf.util.basic import print_graph_output
    print_graph_output(out_lens)
    assert out_lens is in_data.size_placeholder[0]


def test_batch_norm_args():
  from returnn.tf.layers.basic import BatchNormLayer
  from returnn.util.basic import getargspec
  batch_norm_args = getargspec(BatchNormLayer.batch_norm).args[2:]  # drop self and data
  layer_args = getargspec(BatchNormLayer.__init__).args[2:]  # drop self and in_dim
  assert batch_norm_args == layer_args  # different arguments in BatchNormLayer and LayerBase.batch_norm()


def test_pickle_dim_tags():
  # Test for pickling net dict and extern data.
  # https://github.com/rwth-i6/returnn_common/issues/104
  # What kind of network we pickle here is not really relevant.
  # I took this from some other test.
  # It should just involve some dim tags.
  from returnn.tf.util.data import batch_dim, ImplicitDynSizeDim
  n_batch, n_time, n_ts, n_in, n_out = 2, 3, 6, 7, 11
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("in-feature", dimension=n_in)
  ts_dim = SpatialDim("ts", dimension=n_ts)
  out_dim = FeatureDim("out-feature", dimension=n_out)
  config = Config({
    "extern_data": {"data": {"dim_tags": [batch_dim, time_dim, feat_dim]}},
    "network": {
      "t": {
        "class": "range_in_axis", "axis": "t", "from": "data",
        "out_shape": {time_dim, ImplicitDynSizeDim(batch_dim)}},  # (T,)
      "range": {"class": "range", "limit": n_ts, "out_spatial_dim": ts_dim},  # (Ts,)
      "add_t": {
        "class": "combine", "kind": "add", "from": ["t", "range"],
        "out_shape": {time_dim, ts_dim, ImplicitDynSizeDim(batch_dim)}},  # (T,Ts)
      "t_rel_var": {"class": "variable", "shape": (ts_dim, out_dim), "init": "glorot_uniform"},  # (Ts,D)
      "output": {
        "class": "scatter_nd", "from": "t_rel_var", "position": "add_t", "position_axis": ts_dim,
        "out_spatial_dim": time_dim, "filter_invalid_indices": True}  # (T,T,D)
    }
  })

  # First test as-is.
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    network.initialize_params(session)
    fetches = network.get_fetches_dict()
    out_layer = network.get_default_output_layer()
    print("out layer:", out_layer)
    session.run((out_layer.output.placeholder, fetches), feed_dict=make_feed_dict(network.extern_data))

  def _debug_dump(tuple_path, obj):
    print("%s: %s" % (tuple_path, type(obj)))
    assert isinstance(obj, (Dim, str, set, str, bool, int))

  from tensorflow.python.util import nest
  nest.map_structure_with_tuple_paths(_debug_dump, config.typed_dict)

  # Now pickle, unpickle and test again.
  # Use pickle._dumps for easier debugging.
  # noinspection PyUnresolvedReferences
  from pickle import _dumps as dumps
  s = dumps(config.typed_dict)
  import pickle
  config_dict = pickle.loads(s)
  new_dim_tags = config_dict["extern_data"]["data"]["dim_tags"]
  new_batch, new_time, new_feat = new_dim_tags
  assert isinstance(new_batch, Dim) and new_batch == batch_dim and new_batch.is_batch_dim()
  assert isinstance(new_time, Dim) and new_time.is_spatial_dim() and new_time.dimension is None
  assert isinstance(new_feat, Dim) and new_feat.is_feature_dim() and new_feat.dimension == n_in
  config = Config(config_dict)
  with make_scope() as session:
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_dict["network"])
    network.initialize_params(session)
    fetches = network.get_fetches_dict()
    out_layer = network.get_default_output_layer()
    print("out layer:", out_layer)
    session.run((out_layer.output.placeholder, fetches), feed_dict=make_feed_dict(network.extern_data))


def test_contrastive_loss():
  from returnn.tf.util.data import batch_dim
  masked_time_dim = SpatialDim("masked_time")
  input_dim = FeatureDim("input", 10)
  enc_feat_dim = FeatureDim("encoder_dim", 20)
  project_dim = FeatureDim("project_dim", 15)

  dim_neg_samples = SpatialDim("neg_samples", 10)  # 100
  dim_expand = SpatialDim("expand_dim", 1)
  contrastive_loss_temp = 0.1
  contrastive_loss_factor = 1.
  seed = 1

  def _mask(x, axis, pos, max_amount):
    """
    :param tf.Tensor x: (batch,time,feature)
    :param int axis:
    :param tf.Tensor pos: (batch,)
    :param int max_amount: inclusive
    :return: (batch,dim)
    """
    from returnn.tf.compat import v1 as tf
    n_batch = tf.shape(x)[0]
    dim = tf.shape(x)[axis]
    amount = tf.random_uniform(shape=(n_batch,), minval=1, maxval=max_amount + 1, seed=seed, dtype=tf.int32)
    pos2 = tf.minimum(pos + amount, dim)
    idxs = tf.expand_dims(tf.range(0, dim), 0)  # (1,dim)
    pos_bc = tf.expand_dims(pos, 1)  # (batch,1)
    pos2_bc = tf.expand_dims(pos2, 1)  # (batch,1)
    cond = tf.logical_and(tf.greater_equal(idxs, pos_bc), tf.less(idxs, pos2_bc))  # (batch,dim)
    return cond

  def _random_mask(x, axis, min_num, max_num, max_dims):
    """
    :param tf.Tensor x: (batch,time,feature)
    :param int axis:
    :param int|tf.Tensor min_num:
    :param int|tf.Tensor max_num: inclusive
    :param int max_dims: inclusive
    :return: (batch,dim)
    """
    from returnn.tf.compat import v1 as tf
    n_batch = tf.shape(x)[0]
    num = tf.random_uniform(shape=(n_batch,), minval=min_num, maxval=max_num + 1, seed=seed, dtype=tf.int32)
    # https://github.com/tensorflow/tensorflow/issues/9260
    # https://timvieira.github.io/blog/post/2014/08/01/gumbel-max-trick-and-weighted-reservoir-sampling/
    z = -tf.log(-tf.log(tf.random_uniform((n_batch, tf.shape(x)[axis]), 0, 1, seed=seed)))
    _, indices = tf.nn.top_k(z, tf.reduce_max(num))
    return _mask(x, axis=axis, pos=indices[:, 0], max_amount=max_dims)

  def _get_mask_eval_layer(source, **_kwargs):
    data = source(0, as_data=True, auto_convert=False)
    assert (data.batch_dim_axis, data.time_dim_axis) == (0, 1)
    x = data.placeholder
    mask = _random_mask(x, axis=data.time_dim_axis, min_num=1, max_num=1, max_dims=3)
    return mask

  net_dict = {
    "input": {"class": "linear", "from": "data", "activation": "relu", "out_dim": input_dim},

    # True -> should be masked out by mask embed vector, False -> keep
    "input_mask": {
      "class": "eval", "from": "input",
      "eval": _get_mask_eval_layer,
      "out_type": {"dtype": "bool", "shape": (None,)}
    },  # [B,T]

    "mask_emb": {"class": "variable", "shape": [input_dim], "init": "RandomUniform(0., 1.)"},

    "input_masked": {
      "class": "switch", "condition": "input_mask", "true_from": "mask_emb", "false_from": "input"},  # [B,T,F]

    # For this loss, the model must be able to look at context, otherwise it does not make sense.
    # https://github.com/rwth-i6/returnn/pull/918
    "conv": {
      "class": "conv", "from": "input_masked",
      "filter_size": [5], "padding": "same",
      "n_out": 10, "activation": "relu"},

    "encoder": {"class": "linear", "from": "conv", "out_dim": enc_feat_dim},  # [B,T,F]

    "contrastive_loss": {
      "class": "subnetwork", "from": [],
      "subnetwork": {
        "enc_masked_frames_": {"class": "masked_computation", "mask": "base:input_mask", "from": "base:encoder",
                               "unit": {"class": "copy", "from": "data"}},  # [B, T_M, F]
        "enc_masked_frames": {
          "class": "reinterpret_data", "from": "enc_masked_frames_", "set_dim_tags": {"T": masked_time_dim}},
        "c": {"class": "linear", "from": "enc_masked_frames", "out_dim": project_dim},

        # We take the non-masked input of the masked frames -> q_t in the paper.
        "input_masked_frames": {"class": "masked_computation", "mask": "base:input_mask", "from": "base:input",
                                "out_spatial_dim": masked_time_dim,
                                "unit": {"class": "copy", "from": "data"}},  # [B, T_M, F]
        "q": {"class": "linear", "from": "input_masked_frames", "out_dim": project_dim},
        "q_len": {"class": "length", "from": "input_masked_frames", "axis": "T"},  # [B]

        # Candidate samples
        "q_samples_rand_indices__": {
          "class": "rand_int", "maxval": 2 ** 30,
          "from": "input_masked_frames",  # only for masked_time_dim
          "shape": [batch_dim, masked_time_dim, dim_neg_samples]},  # [B, T_M, K] -> 0..BIG
        "q_samples_rand_indices_": {
          "class": "eval", "from": ["q_samples_rand_indices__", "q_len"],
          "eval": "source(0) % tf.maximum(source(1) - 1, 1)"},  # [B, T_M, K] -> 0..T_M-1
        "_range": {"class": "range_in_axis", "from": "input_masked_frames", "axis": masked_time_dim},  # [T_M]
        "_range_ge_indices": {
          "class": "compare", "kind": "greater_equal", "from": ["q_samples_rand_indices_", "_range"]},  # [B, T_M, K]
        "_indices_offsets": {"class": "switch", "condition": "_range_ge_indices", "true_from": 1, "false_from": 0},
        "q_samples_rand_indices": {
          "class": "combine", "kind": "add", "from": ["q_samples_rand_indices_", "_indices_offsets"]},  # [B, T_M, K]
        "q_sampled_frames": {
          "class": "gather", "from": "q",
          "position": "q_samples_rand_indices", "axis": masked_time_dim},  # [B, T_M, K, F]
        "q_expand": {"class": "expand_dims", "axis": "spatial", "dim": dim_expand, "from": "q"},  # [B, T_M ,1, F]
        "Q": {"class": "concat", "from": [
          ("q_expand", dim_expand), ("q_sampled_frames", dim_neg_samples)]},  # [B, T_M, K+1, F]

        # Cosine similarity between sampled frames and masked encoder frames
        "cos_similarity": {
          "class": "subnetwork", "from": ["Q", "c"], "concat_sources": False,
          "subnetwork": {
            # [B_M, K+1, F] * [B_M, F] -> [B_M, K+1]
            "dot": {"class": "dot", "from": ["data:0", "data:1"], "reduce": project_dim},
            "norm_a_sq_": {"class": "eval", "from": "data:0", "eval": "source(0) ** 2"},
            "norm_a_sq": {"class": "reduce", "mode": "sum", "from": "norm_a_sq_", "axes": project_dim},  # [B, T_M, K+1]
            "norm_b_sq_": {"class": "eval", "from": "data:1", "eval": "source(0) ** 2"},
            "norm_b_sq": {"class": "reduce", "mode": "sum", "from": "norm_b_sq_", "axes": project_dim},  # [B, T_M]
            "output": {
              "class": "eval", "from": ["dot", "norm_a_sq", "norm_b_sq"],
              "eval": "source(0) * tf.minimum(tf.math.rsqrt(source(1) * source(2)), 1./1e-8)"},  # [B, T_M, K+1]
          },
        },

        # The contrastive loss is the negative log-likelihood of the softmax of the cosine similarity
        "log_sm_cos_sim": {
          "class": "softmax_over_spatial", "from": "cos_similarity", "axis": dim_expand + dim_neg_samples,
          "log_space": True, "energy_factor": contrastive_loss_temp},  # [B, T_M, K+1]
        "log_likelihood": {
          "class": "gather", "from": "log_sm_cos_sim", "axis": dim_expand + dim_neg_samples, "position": 0},  # [B, T_M]
        "neg_los_likelihood": {"class": "eval", "from": "log_likelihood", "eval": "-source(0)"},  # [B, T_M]
        "output": {"class": "copy", "from": "neg_los_likelihood"},
      },
      "loss": "as_is", "loss_scale": contrastive_loss_factor,
    },

    "output": {"class": "softmax", "from": "encoder", "loss": "ce"}
  }

  with make_scope() as session:
    config = Config({
      "extern_data": {"data": {"dim": 7}, "classes": {"dim": 3, "sparse": True}},
      "debug_print_layer_output": True, "debug_print_layer_output_shape": True,
    })
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict(net_dict)
    loss = net.get_total_loss()

    tf_compat.v1.set_random_seed(1)
    net.initialize_params(session)
    loss_v = session.run(loss, feed_dict=make_feed_dict(net.extern_data, same_time=True))
    print("loss:", loss_v)
    assert numpy.isfinite(loss_v)


if __name__ == "__main__":
  try:
    better_exchook.install()
    if len(sys.argv) <= 1:
      for k, v in sorted(globals().items()):
        if k.startswith("test_"):
          print("-" * 40)
          print("Executing: %s" % k)
          try:
            v()
          except unittest.SkipTest as exc:
            print("SkipTest:", exc)
          print("-" * 40)
      print("Finished all tests.")
    else:
      assert len(sys.argv) >= 2
      for arg in sys.argv[1:]:
        print("Executing: %s" % arg)
        if arg in globals():
          globals()[arg]()  # assume function and execute
        else:
          eval(arg)  # assume Python code and execute
  finally:
    import threading
    #if len(list(threading.enumerate())) > 1:
    #  print("Warning, more than one thread at exit:")
    #  better_exchook.dump_all_thread_tracebacks()