bool IsOpLegalizedWithMlir(Operation& op);

// Given the type ID, check if it's legalized with MLIR.
bool IsTypeLegalizedWithMlir(const TypeID& type_id);

// Returns true if the op is considered a dynamic padder op.
bool IsDynamicPadderOp(const TypeID& type_id);

// Returns True if this op has a Tf2XLA fallback. Currently, this is not the
// inverse of the !IsOpLegalizedWithMlir, but it should be.
bool HasTf2XlaFallback(const TypeID& type_id);

  EXPECT_FALSE(IsTypeLegalizedWithMlir(TypeID::get<TF::AllOp>()));
}

TEST_F(LegalizationOpConfigTest, ChecksDynamicPadderOps) {
  EXPECT_TRUE(
      IsDynamicPadderOp(TypeID::get<TF::XlaSetDynamicDimensionSizeOp>()));
  EXPECT_FALSE(IsDynamicPadderOp(TypeID::get<TF::ConstOp>()));
}

// This test is kind of odd. We go through all the Tensorflow types and check

  Operation* func_op = getOperation();

  func_op->walk([&](Operation* op) {
    auto abstractOp = op->getRegisteredInfo();
    if (!abstractOp) return WalkResult::advance();

    if (mlir::mhlo::IsDynamicPadderOp(abstractOp->getTypeID())) {
      has_dynamic_op = true;
      dynamism_op_counter->GetCell(op->getName().getStringRef().str())
          ->IncrementBy(1);
    }

    return WalkResult::advance();
  });

  return IsOpTypeAllowedTf2XlaFallback(type_id);
}

bool IsOpAllowedTf2xlaPreferred(const TypeID& type_id) {
  return IsOpTypeAllowedTf2XlaPreferred(type_id);
}

bool IsDynamicPadderOp(const TypeID& type_id) {
  return DynamicTensorflowOps().contains(type_id);
}

}  // namespace mhlo