message: [tf.data] Fixing disable intra-op parallelism rewrite.
diff: @@ -89,9 +89,19 @@ Status DisableIntraOpParallelism::OptimizeAndCollectStats(
  // `max_intra_op_parallelism` input
  *insert_node.mutable_input()->Add() = max_parallelism_value->name();
  - for (const auto& attr_name : {"output_types", "output_shapes"}) {
  - if (last_node->attr().contains(attr_name)) {
  - graph_utils::CopyAttribute(attr_name, *last_node, &insert_node);
  + // Set `output_types` and `output_shapes` attributes by copying the relevant
  + // attrs from the input node. If we fail to set the attributes, we abort the
  + // rewrite.
  + for (auto attr : {"output_shapes", "output_types"}) {
  + if (last_node->attr().find(attr) != last_node->attr().end()) {
  + graph_utils::CopyAttribute(attr, *last_node, &insert_node);
  + } else {
  + auto inferred_attr = strings::StrCat("T", attr);
  + if (last_node->attr().find(inferred_attr) != last_node->attr().end()) {
  + graph_utils::CopyAttribute(inferred_attr, *last_node, &insert_node);
  + } else {
  + return Status::OK();
  + }
  }
  }
  