// module level.
  for (auto func : module_op.getOps<func::FuncOp>()) {
    if (failed(applyPatternsAndFoldGreedily(func, frozen_patterns))) {
      func.emitError() << "quant-propagate-quantize-type failed.";
      signalPassFailure();
    }
  }
}

}  // namespace

// Creates an instance of the TensorFlow dialect PropagateQuantizeType pass.

    }

    Value input = tf_op.getInputs();
    int quant_dim = -1;
    auto input_type = mlir::cast<ShapedType>(input.getType());
    if (PerAxis) {
      if (!input_type.hasRank()) {
        tf_op.emitError("The input should have known rank for per-channel op.");
        return failure();
      }
      // This is a special case that the quant_dim is the last dimensions.
      quant_dim = input_type.getRank() - 1;
    }

      << ": " << num_part_ops << " ops delegated out of " << num_all_ops
      << " ops with " << num_partitions_after
      << " partitions (originally: " << num_partitions_before << ")";
  if (!scheduled_everything) {
    emitError(func.getLoc(), func.getName())
        << ": " << num_unscheduled_ops << " operations couldn't be scheduled";
  }
  return scheduled_everything;
}

class PartitionedTopologicalSortPass

                                 &send_gradient_op_map)))
    return signalPassFailure();

  if (enqueue_op_map.size() != recv_activation_op_map.size()) {
    func_op.emitError() << "expects the number of embedding enqueue ops to "
                           "match the number of '"
                        << TF::RecvTPUEmbeddingActivationsOp::getOperationName()
                        << "' ops";

  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<RewriteNchwConvolutionToNhwc>(&ctx);

  if (failed(applyPatternsAndFoldGreedily(func_op, std::move(patterns)))) {
    func_op.emitError() << "Failed to run NchwConvolutionToNhwcPass.";
    signalPassFailure();
  }
}

  RewritePatternSet patterns(ctx);
  func::FuncOp func = getOperation();

  patterns.add<AddQuantizationUnitLoc>(ctx);
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    func.emitError() << "quant-add-quantization-unit-loc pattern "
                        "conversion did not converge.";
    signalPassFailure();
  }
}

}  // namespace

  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<FoldTransposedConstantOp>(&ctx);
  if (failed(applyPatternsAndFoldGreedily(func_op, std::move(patterns)))) {
    func_op.emitError("Failed to fold constant->transpose pattern.");
    signalPassFailure();
  }
}

  FrozenRewritePatternSet frozen_patterns(std::move(patterns));
  for (auto func : module.getOps<func::FuncOp>()) {
    if (failed(applyPatternsAndFoldGreedily(func, frozen_patterns))) {
      func.emitError()
          << "quant-lift-quantizable-spots-as-functions-drq failed.";
      signalPassFailure();
    }
  }
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>

    // island wraps a single op.
    Operation& wrapped_op = island.GetBody().front();
    TF::ParallelExecutionIdPairs id_pairs;
    if (failed(TF::ParseParallelExecutionIds(&wrapped_op, id_pairs))) {
      wrapped_op.emitError()
          << "Malformed " << TF::kParallelExecAnnotation << " attribute";
      return failure();
    }
    if (id_pairs.empty()) continue;

    TF::ParallelIdsMap& ids_map = op_to_parallel_ids_map[island];

    if (exportedNames.empty()) {
      name = "global_ml_" + globalTensor.getSymName().str();
    } else if (exportedNames.size() == 1) {
      name = exportedNames[0].str();
    } else {
      return globalTensor.emitError()
             << "Multiple exported names for global tensor not supported yet";
    }
    Attribute initial_value;
    if (globalTensor.getValue()) {
      initial_value = *globalTensor.getValue();
    } else {