ReduceLayer over B, T, F

[vieting]

I have a setup that does reduce mean over (B, T, F) which are all axes in my case. This commit 0e3fe3b breaks the setup with

  File ".../returnn/returnn/tf/layers/basic.py", line 6220, in ReduceLayer.reduce
    line: x.dim_tags[axis].dyn_size_ext
    locals:
      x = <local> Data{'data', [B,T|'time:var:extern_data:data'[B],F|F'feature:data'(4)]}
      x.dim_tags = <local> (Dim{B}, Dim{'time:var:extern_data:data'[B]}, Dim{F'feature:data'(4)})
      axis = <local> 1
      dyn_size_ext = <not found>
  File ".../returnn/returnn/tf/util/data.py", line 3808, in Data.copy_compatible_to
    line: raise ValueError("copy_compatible_to: self %r has batch-dim, but target data %r has not" % (self, data))
    locals:
      ValueError = <builtin> <class 'ValueError'>
      self = <local> Data{'data_dim0_size', [B], dtype='int32'}
      data = <local> Data{'data', []}
ValueError: copy_compatible_to: self Data{'data_dim0_size', [B], dtype='int32'} has batch-dim, but target data Data{'data', []} has not

A test case to reproduce the error:

def test_ReduceLayer_mean_btf():
  net_dict = {
    "output": {"class": "reduce", "mode": "mean", "from": "data", "axis": ["B", "T", "F"]}
  }
  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
    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]
    for b in range(n_batch):
      in_v[b, seq_len[b]:, :] = numpy.nan
    numpy.testing.assert_almost_equal(out_v, numpy.nanmean(in_v))

[vieting]

There is a special case

if x.batch_dim_axis in axes and x.time_dim_axis in axes and len(axes) == 2:

which just flattens batch and time for the reduce op. I wonder why it also checks for len(axes) == 2 here. Is that necessary? In my case, flattening batch and time and then reducing over the remaining axes works perfectly fine. Could we maybe remove it or at least relax it to also apply if we reduce over all axes of the tensor?

[albertz]

I wonder why it also checks for len(axes) == 2 here. Is that necessary?

Yes, it's necessary, because otherwise it would miss further checks and handling for potential other dynamic dims.

But also, this is a special case to (maybe) make it faster, and the more generic code should in any case handle all possible cases correct.

So, why does it actually fail? I realize that the code is anyway not really correct for reduce_mean when there are multiple axes. I think it's wrong to just multiply the correction factors like this. Or is it? But also the dyn_size_ext.copy_compatible_to doesn't seem right. But I'm not sure now how to handle it. reduce_sum over dyn_size_ext in case there are more dims than in out_data?

[albertz]

Ok I think I got it. Should be fixed.

[vieting]

Yes, fixed my actual issue as well. Thank you 👍