return WalkResult::interrupt();
      }

      return WalkResult::advance();
    });

    if (walk_result.wasInterrupted()) {
      return tsl::errors::Internal("Could not legalize all ops");
    }

    return absl::OkStatus();
  }

  mlir::func::FuncOp GetMainFunc() {
    func::FuncOp main_func = module_->lookupSymbol<mlir::func::FuncOp>("main");
    EXPECT_TRUE(main_func);

    "post-training-quantization",
    llvm::cl::desc("Enable post_training_quantization."),
    llvm::cl::init(false));

// NOLINTNEXTLINE
opt<bool> legalize_custom_tensor_list_ops(
    "legalize-custom-tensor-list-ops",
    llvm::cl::desc("Convert \"tf.TensorList*\" ops to \"tfl.custom_op\""
                   "if they can all be supported."),
    llvm::cl::init(false));

// NOLINTNEXTLINE

    // Do not append a UnitAttr for the "token" operand here to avoid
    // compilation failure when exporting the "layouts" attribute to a graph
    // node. Instead, add the UnitAttr during LegalizeTF pass.
    op->setAttr("layouts", layout.value());

    return mlir::WalkResult::advance();
  });
  return !res.wasInterrupted();
}

            tasks.create("foo", Foo)
            tasks.register("bar", Foo) {
                assert false : "This should not be realized"
            }
            tasks.create("builtInTask", Copy)
            tasks.register("defaultTask") {
                assert false : "This should not be realized"
            }

            def schema = tasks.collectionSchema.elements.collectEntries { e ->

 public:
  explicit TflToStablehloPass() : PassWrapper() {}
  StringRef getArgument() const final { return "tfl-parse-stablehlo-ops"; }
  StringRef getDescription() const final {
    return "This pass will legalize TFLite custom Ops to StableHLO ops.";
  }

 private:
  void runOnOperation() override;

  void getDependentDialects(DialectRegistry& registry) const override {

                                        PatternRewriter& rewriter,
                                        const std::string& device_type) {
  Tf2XlaRewriter tf2xla_rewriter(op, rewriter, device_type);
  return tf2xla_rewriter.LegalizeOp();
}

Tf2XlaRewriter::Tf2XlaRewriter(Operation* op, PatternRewriter& rewriter,
                               const std::string& device_type)
    : op_(op),
      device_type_(device_type),

bool GetDummyParams(OpBuilder& builder, Value val_bcast, Attribute& zero,
                    DenseIntElementsAttr& shape) {
  Type type = val_bcast.getType();
  Type elem_type = getElementTypeOrSelf(type);
  // Xla's all_reduce legalizer bitcasts to 32 bits, so only
  // element types size <= 4 bytes are supported.
  if (elem_type.isBF16() || elem_type.isF16() || elem_type.isTF32() ||
      elem_type.isF32()) {
    zero = builder.getFloatAttr(elem_type, 0);

import org.gradle.test.preconditions.IntegTestPreconditions

class DeferredTaskFailureIntegrationTest extends AbstractDeferredTaskDefinitionIntegrationTest {
    def "reports failure in task constructor when task realized"() {
        createDirs("child")
        settingsFile << """
            include "child"
        """
        file("child/build.gradle") << """
            class Broken extends DefaultTask {
                Broken() {

                         OperationPass<ModuleOp>> {
 public:
  StringRef getArgument() const final {
    return "tf-xla-call-module-op-to-stablehlo-pass";
  }
  StringRef getDescription() const final {
    return "Legalize TF_XlaCallModule Op to stablehlo";
  }
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<stablehlo::StablehloDialect, vhlo::VhloDialect,

  // `lower_static_tensor_list` pass.
  return element_type->isF32() || element_type->isInteger(64) ||
         element_type->isInteger(32) || element_type->isInteger(1);
}

// Only legalize TensorFlow TensorList ops if all TensorList ops are supported
// natively.
class LegalizeTensorListPass
    : public impl::LegalizeTensorListPassBase<LegalizeTensorListPass> {
 public: