/*tf2xla_fallback_device_type=*/device_type, enable_op_fallback));
  }

  // This has to run after legalization.
  pm.addNestedPass<mlir::func::FuncOp>(
      mlir::mhlo::CreateInfeedsOpsXlaAdjustLayoutPass());

  if (lower_to_xla_hlo) {
    // This has to run after legalization to delete non legal but dead ops.
    // This must run before Shape Inference.

  let description = [{
      Pass which analyzes the variables in the graph and add an attribute whether
      variables should be legalized to TFLite native ones.
      This pass needs to run post TF->TFL legalization and before variable
      legalization.
  }];
  let constructor = "CreateAnalyzeVariablesPass()";
  let dependentDialects = ["TFL::TensorFlowLiteDialect"];
}

  // broadcasting support. This needs to be run immediately after HLO->TF
  // legalization; otherwise other passes like `ConvertTFBroadcastTo` will
  // constant fold the newly generated TF broadcast ops and materialize the
  // weights.
  pass_manager.addNestedPass<mlir::func::FuncOp>(
      mlir::TF::CreateBroadcastFoldPass());

  // Canonicalization after TF legalization.
  pass_manager.addNestedPass<mlir::func::FuncOp>(

  pass_manager.printAsTextualPipeline(raw_stream);

  EXPECT_THAT(pass_description, HasSubstr(kInlinePass));
}

TEST(LegalizeMlirTest, LegalizesModuleWithDynamicShape) {
  constexpr char legalization[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

    %1 = "tf.Reshape"(%arg1, %0) : (tensor<19x10xf32>, tensor<2xi64>) -> tensor<10x19xf32>
    func.return %1 : tensor<10x19xf32>
  }
}

// Tests that foldable ops are constant-folded to enable legalization of ops
// that require compile time constant operand.
// "tf.Shape" can only be folded away after shape inference. tf.Reshape can only
// be lowered when tf.Shape is folded into a constant.

// CHECK-LABEL: HloModule main

}

// Adds the list of ops that are supported on TPU through constant folding which
// may depend on the inputs shapes not known at this point. Such ops may not
// have any legalization or canonicalization patterns but shouldn't be marked
// for outside compilation.
//
// TODO(b/177523289): Remove manual handling once we support constant folding

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeVariablesPass)

  void runOnOperation() override {
    auto module = getOperation();
    // If TFLite variable legalization is not allowed, then we skip this pass.
    if (auto legalize_tfl_variables_attr =
            module->getAttr(kLegalizeTflVariables)) {
      if (!mlir::cast<BoolAttr>(legalize_tfl_variables_attr).getValue()) return;
    }

    expressions.push_back(expr);

    if (!tensorflow::DataTypeCanUseMemcpy(expr.dtype())) {
      return op_->emitRemark()
             << "skipping legalization due to unsupported type "
             << operand.getType();
    }

    auto shape_or = expr.GetShape();
    if (!shape_or.ok()) {
      return op_->emitRemark()

//    {
//      ...
//    }
class TPUEmbeddingOpsRegistry {
 public:
  // Add the op to the registry.
  //
  // Adding an op here will allow use old bridge legalization from the MLIR
  // bridge with the use of fallback mechanism. Therefore, addition of any op
  // here must have a python test with MLIR bridge enabled to verify that the
  // fallback works correctly.

  pm.addNestedPass<func::FuncOp>(mlir::createCanonicalizerPass());
  pm.addPass(mhlo::createHloLegalizeToStablehloPass());
}

void AddStablehloQuantToIntPasses(OpPassManager& pm) {
  // StableHLO -> MHLO legalization.
  pm.addPass(mhlo::createStablehloLegalizeToHloPass());
  pm.addNestedPass<func::FuncOp>(mhlo::createMhloQuantLegalizeToIntPass());
  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());