auto arg_type = mlir::cast<RankedTensorType>(arg.getType());
      assert(arg_type.getRank() == 0);
      llvm::ArrayRef<TensorType> underlying_type =
          mlir::cast<TF::ResourceType>(arg_type.getElementType()).getSubtypes();

      // If the arg type already matches the global_tensor type, we don't need
      // to do anything.
      if (!underlying_type.empty() &&

  int64_t output_size = InferTensorSize(context, value_type);

  int64_t cost = kLookupTableFindCostScale * output_size;

  if (mlir::isa<mlir::TF::StringType>(key_type.getElementType()))
    cost *= kLookupTableFindStringKeyCostScale;

  return cost;
}

// The cost function for tf.GatherV2.
int64_t InferGatherV2Cost(const CostContext& context, mlir::TF::GatherV2Op op) {

                split_dimension, num_split));
      }

      shape[split_dimension] = shape[split_dimension] / num_split;
      output_type =
          mlir::RankedTensorType::get(shape, input_type.getElementType());
    }
  } else {
    output_type = input_type;
  }

  // Creates a split op that splits |src_input| along |split_dimension|.
  llvm::SmallVector<mlir::Type, 4> output_types(num_split, output_type);

        TypedProviderOfSet(Class<T> type, Set<T> value) {
            this.type = type
            this.value = value
        }

        @Override
        Class<? extends T> getElementType() {
            return type
        }

        @Override
        protected Value<? extends Set<T>> calculateOwnValue(ValueConsumer consumer) {
            return Value.of(value)
        }

  // require "raising" later.
  for (const Type type : inst->getResultTypes()) {
    if (const TensorType tensor_type = mlir::dyn_cast<TensorType>(type)) {
      if (mlir::isa<VariantType>(tensor_type.getElementType())) {
        return false;
      }
    }
  }

  // Operations that execute function calls shouldn't be constant folded.
  if (llvm::isa<TF::WhileOp, TF::CaseOp, TF::IfOp, CallOpInterface>(inst)) {

    // the operand can be tensor of type T or T_REF.
    bool is_output_ref = mlir::isa<tf_type::TensorFlowRefType>(
        output_tensor_type.getElementType());
    if (is_output_ref && !mlir::isa<tf_type::TensorFlowRefType>(
                             operand0_tensor_type.getElementType())) {
      return switchn.emitOpError()
             << "expects same operand and output element type but got "

      ShapedType value_shaped_type = mlir::dyn_cast_or_null<ShapedType>(
          value_for_deducing_op_type.getType());
      if (value_shaped_type != nullptr) {
        auto operand_or_result = value_shaped_type.getElementType();
        std::string dtype;

        TypeSwitch<Type>(operand_or_result)
            .Case<IntegerType>([&](Type) {
              dtype =

namespace mlir {
namespace TF {
namespace {

bool IsResourceType(Type type) {
  if (auto tensor_type = mlir::dyn_cast<TensorType>(type)) {
    return mlir::isa<TF::ResourceType>(tensor_type.getElementType());
  }
  return false;
}

bool IsResource(Value value) { return IsResourceType(value.getType()); }

// Helper that returns the FuncOp that is the SessionInit function which

  if (!broadcasted_shapes_or.has_value()) return;
  const auto broadcasted_input_type = RankedTensorType::get(
      broadcasted_shapes_or->first, input_type.getElementType());
  const auto broadcasted_weight_type = RankedTensorType::get(
      broadcasted_shapes_or->second, weight_type.getElementType());

  if (broadcasted_input_type.hasStaticShape() &&
      broadcasted_weight_type.hasStaticShape()) {
    input = builder.create<TF::BroadcastToOp>(

bool IsResourceTensor(Value value) {
  const auto tensor_type = mlir::dyn_cast<TensorType>(value.getType());
  return tensor_type &&
         mlir::isa<mlir::TF::ResourceType>(tensor_type.getElementType());
}

// The default criterion for operations being considered as causing or being
// dependent on side effects. Reflects the current runtime logic; see below.
bool OpHasSideEffects(Operation *op) {