message: Re-enable the 5d fused batch norm
diff: @@ -235,25 +235,27 @@ class BatchNormalizationBase(Layer):
  """Raises a ValueError if fused implementation cannot be used.
  In addition to the checks done in this function, the input tensors rank must
  - be 4. The input rank check can only be done once the input shape is known.
  + be 4 or 5. The input rank check can only be done once the input shape is
  + known.
  """
  # Note the ValueErrors in this function are caught and not reraised in
  # _fused_can_be_used(). No other exception besides ValueError should be
  # raised here.
  # Currently fused batch norm doesn't support renorm. It also only supports a
  - # channel dimension on axis 1 or 3, when no virtual batch size or adjustment
  - # is used.
  + # channel dimension on axis 1 or 3 (rank=4) / 1 or 4 (rank5), when no
  + # virtual batch size or adjustment is used.
  if self.renorm:
  raise ValueError('Passing both `fused=True` and `renorm=True` is '
  'not supported')
  axis = [self.axis] if isinstance(self.axis, int) else self.axis
  - # Axis -3 is equivalent to 1, and axis -1 is equivalent to 3, because the
  - # input rank is required to be 4 (which is checked later).
  - # TODO(b/173253101): Once the input rank can be 5, update this check.
  - if len(axis) > 1 or axis[0] not in (-3, -1, 1, 3):
  + # Axis -3 is equivalent to 1, and axis -1 is equivalent to 3, when the
  + # input rank is 4. Similarly, the valid axis is -4, -1, 1, 4 when the rank
  + # is 5. The combination of ranks and axes will be checked later.
  + if len(axis) > 1 or axis[0] not in (-4, -3, -1, 1, 3, 4):
  raise ValueError('Passing `fused=True` is only supported when axis is 1 '
  - 'or 3. Got axis %s' % (axis,))
  + 'or 3 for input rank = 4 or 1 or 4 for input rank = 5. '
  + 'Got axis %s' % (axis,))
  if self.virtual_batch_size is not None:
  raise ValueError('Passing `fused=True` is not supported when '
  '`virtual_batch_size` is specified.')
  @@ -334,19 +336,16 @@ class BatchNormalizationBase(Layer):
  # TODO(yaozhang): if input is not 4D, reshape it to 4D and reshape the
  # output back to its original shape accordingly.
  if self._USE_V2_BEHAVIOR:
  - # TODO(b/173253101): Using fused in the 5D case is currently disabled
  - # due to a regression on UNet, so it is only currently only supported in
  - # the 4D case.
  if self.fused is None:
  - self.fused = ndims == 4
  - elif self.fused and ndims != 4:
  + self.fused = ndims in (4, 5)
  + elif self.fused and ndims not in (4, 5):
  raise ValueError('Batch normalization layers with `fused=True` only '
  'support 4D or 5D input tensors. '
  'Received tensor with shape: %s' %
  (tuple(input_shape),))
  else:
  assert self.fused is not None
  - self.fused = (ndims == 4 and self._fused_can_be_used())
  + self.fused = (ndims in (4, 5) and self._fused_can_be_used())
  # TODO(chrisying): fused batch norm is currently not supported for
  # multi-axis batch norm and by extension virtual batches. In some cases,
  # it might be possible to use fused batch norm but would require reshaping
  