getOperation(), max_iterations_, /*ops_to_skip=*/{}, input_shapes_);
    if (failed(failure_or_converged)) return signalPassFailure();
    if (!failure_or_converged.value()) {
      getOperation().emitError()
          << "shape inference pass did not reach convergence after "
          << max_iterations_;
      return signalPassFailure();
    }
  }

 private:
  ArrayRef<ArrayRef<int64_t>> input_shapes_;
};
}  // namespace

      module_in, confs, &graph, &flib_def, &control_ret_nodes);
  if (!status.ok()) {
    mlir::emitError(mlir::UnknownLoc::get(&ctx)) << status.message();
    return signalPassFailure();
  }

  // Run each of the passes that were selected.
  GraphConstructorOptions opts;
  opts.allow_internal_ops = true;
  opts.expect_device_spec = false;

  GraphOptimizationPassOptions options;

    if (failed(ConvertTFTextAPI(func, api_name, attr))) {
      return signalPassFailure();
    }
  } else if (api_name == kCustomSSDPostprocessing) {
    ConvertSSDPostProcessFunc convert_ssd_postprocess(func, attr);
    if (failed(convert_ssd_postprocess.VerifySignature())) return;
    if (failed(convert_ssd_postprocess.RewriteFunc())) {
      return signalPassFailure();
    }
  } else if (api_name == kCustomMaxUnpooling) {

  DataFlowSolver solver;
  solver.load<dataflow::DeadCodeAnalysis>();
  solver.load<dataflow::SparseConstantPropagation>();
  TF::LoadResourceDataflowAnalysis(solver);
  if (failed(solver.initializeAndRun(module))) return signalPassFailure();

  DenseSet<GlobalTensorOp> remaining_global_tensor_ops;
  {
    auto ops = module.getOps<GlobalTensorOp>();
    remaining_global_tensor_ops.insert(ops.begin(), ops.end());
  }

      op->emitOpError() << "embedding ops in two different regions";
      return signalPassFailure();
    }
    if (GetReplicationAttr(op) != replication_attr) {
      op->emitOpError() << "embedding ops with different replication info "
                        << replication_attr << " vs " << GetReplicationAttr(op);
      return signalPassFailure();
    }
  }

  // TODO(bfontain): Check that the region here is the region

  if (entry_funcs.empty()) {
    LOG(WARNING) << "missing entry functions";
  }

  if (HasNoNestedEntryFunctions(entry_funcs, symtab).failed()) {
    return signalPassFailure();
  }

  if (HasTopLevelCompilationMarker(entry_funcs).failed()) {
    return signalPassFailure();
  }
}

}  // namespace

namespace TFDevice {
std::unique_ptr<OperationPass<ModuleOp>> CreateXlaValidateInputsPass() {

      op->emitOpError() << "embedding ops in two different regions";
      return signalPassFailure();
    }
    if (GetReplicationAttr(op) != replication_attr) {
      op->emitOpError() << "embedding ops with different replication info "
                        << replication_attr << " vs " << GetReplicationAttr(op);
      return signalPassFailure();
    }
  }

  // TODO(bfontain): Check that the region here is the region

void RewriteTPUEmbeddingOps::runOnOperation() {
  func::FuncOp func = getOperation();
  if (failed(RunOnRegion(&func.getBody()))) return signalPassFailure();

  func.walk([&](Operation* op) {
    for (Region& region : op->getRegions()) {
      if (failed(RunOnRegion(&region))) return signalPassFailure();
    }
  });
}

}  // anonymous namespace

std::unique_ptr<OperationPass<func::FuncOp>>

    }
    variable_v2s_to_replace.push_back(variable_v2_op);
    return mlir::WalkResult::advance();
  };

  WalkResult walk_res = func.walk(read_only_vars_fn);
  if (walk_res.wasInterrupted()) return signalPassFailure();

  for (VariableV2Op variable_v2_op : variable_v2s_to_replace) {
    builder.setInsertionPoint(variable_v2_op);
    ShapedType shaped_type =
        mlir::cast<ShapedType>(variable_v2_op.getResult().getType());

    return signalPassFailure();
  }

  auto result = getOperation().walk([&tf_dialect](tf_device::LaunchOp launch) {
    if (failed(HoistOpsAndAnnotateWithDevice(tf_dialect, launch)))
      return WalkResult::interrupt();

    return WalkResult::advance();
  });

  if (result.wasInterrupted()) return signalPassFailure();

  if (!legacy_graph_export_) {