The first dim length of the axisStats attribute should equals to the slice size of the input argument when splitted by the axis dimension. PiperOrigin-RevId: 272798042
The first dim length of the axisStats attribute should equals to the slice size of the input argument when splitted by the axis dimension. PiperOrigin-RevId: 272798042