This code is brittle and dangerous. Should be something like:

    axes_to_reduce = range(4)
    del axes_to_reduce[self.output_space.get_batch_axis()]
    ave = total.mean(axis=axes_to_reduce)