The definition of flat should refer to from autograd.util import flatten, but actually refers to a function in the same file with the same name "flatten". This leads to a bug when "flat" is used. The expected_stats should be a 1-D ndarray, but it is a scalar.