Infers the data format for the fused batch norm.

    It uses the axis option to infer this information. Specifically, the
    axis value (0, 1, 2) corresponds to data format NHWC and the
    axis value (0, 2, 3) to data format NCHW.

    Args:
      input_batch: A Tensor of arbitrary dimension.

    Returns:
      A string description of the data format NHWC or NCHW.

    Raises:
      NotImplementedError: for input of dimensionality different from 4.
      ValueError: for axis configuration different from (0, 1, 2) and (0, 2, 3).

def _infer_fused_data_format(self, input_batch):
    """Infers the data format for the fused batch norm.

    It uses the axis option to infer this information. Specifically, the
    axis value (0, 1, 2) corresponds to data format NHWC and the
    axis value (0, 2, 3) to data format NCHW.

    Args:
      input_batch: A Tensor of arbitrary dimension.

    Returns:
      A string description of the data format NHWC or NCHW.

    Raises:
      NotImplementedError: for input of dimensionality different from 4.
      ValueError: for axis configuration different from (0, 1, 2) and (0, 2, 3).
    """
    input_shape = input_batch.get_shape().as_list()
    input_shape_len = len(input_shape)
    if input_shape_len != 4:
      raise NotImplementedError("fused batch norm supports only input with "
                                "4 dimensions, it received input of "
                                "dimensionality {:d}".format(input_shape_len))
    axis = range(input_shape_len)[:-1] if self._axis is None else self._axis
    axis = tuple(axis)
    if axis == (0, 1, 2):
      # Reduce over the last dimension.
      return "NHWC"
    elif axis == (0, 2, 3):
      # Reduce over the second dimension.
      return "NCHW"
    else:
      raise ValueError("Invalid axis option {}. This does not correspond to"
                       " either the NHWC format (0, 1, 2) or the NCHW "
                       "(0, 2, 3).".format(axis))