message: Add layout opt support of 5D batch norm
diff: @@ -670,7 +670,19 @@ Status LayoutSensitiveOpTransposer::UpdateNode(TransposeContext* context,
  Status DefaultLayoutSensitiveOpTransposer::TransposeNode(
  TransposeContext* context, utils::MutableNodeView* node) {
  DCHECK(IsDefaultLayoutSensitiveOp(*node->node()));
  - if (!ShouldProcess(*context, *node) || !IsFanoutPortRankN(*node, 0, 4)) {
  + const auto* output_shape_attr = node->GetAttr(kAttrOutputShape);
  + const auto& shape = output_shape_attr->list().shape(0);
  + const int rank = shape.dim_size();
  + std::string src_format = context->src_format;
  + std::string dst_format = context->dst_format;
  + // Update the format from 4D to 5D layout if necessary.
  + if (rank == 5) {
  + std::string src_format_3d = src_format == "NHWC" ? "NDHWC" : "NCDHW";
  + std::string dst_format_3d = dst_format == "NHWC" ? "NDHWC" : "NCDHW";
  + context->AssignDeviceAndDataFormats(context->target_device, src_format_3d,
  + dst_format_3d);
  + }
  + if (!ShouldProcess(*context, *node) || !IsFanoutPortRankN(*node, 0, rank)) {
  return Status::OK();
  }
  VLOG(3) << "GenericLayoutOptimizer: transforming node '" << node->GetName()
  @@ -679,6 +691,11 @@ Status DefaultLayoutSensitiveOpTransposer::TransposeNode(
  TF_RETURN_IF_ERROR(UpdateNode(context, node));
  TF_RETURN_IF_ERROR(UpdateFaninEdgesWithOp(context, {0}, node, kOpTranspose));
  TF_RETURN_IF_ERROR(UpdateFanoutEdgesWithOp(context, {0}, node, kOpTranspose));
  + // Change back the format from 5D to 4D layout.
  + if (rank == 5) {
  + context->AssignDeviceAndDataFormats(context->target_device, src_format,
  + dst_format);
  + }
  return context->graph_view->GetMutationBuilder()->Apply();
  }
  