SmallVector<int64_t, 4> permutedShape = applyPermutation(
      type.getShape(), isInvert ? invertPermutationVector(perm) : perm);
  return RankedTensorType::get(permutedShape, type.getElementType());
}

static RankedTensorType GetInvertPermutedTensorType(RankedTensorType type,
                                                    ArrayRef<int64_t> perm) {

    if (!const_op) return failure();

    // Only support float tensors.
    auto tensor_type = mlir::dyn_cast_or_null<TensorType>(const_op.getType());
    if (!tensor_type || !tensor_type.getElementType().isF32()) {
      return failure();
    }

    return success(
        mlir::isa_and_nonnull<DenseFPElementsAttr>(const_op.getValue()));
  }

  void rewrite(mlir::stablehlo::TransposeOp op,

    auto tensor =
        llvm::dyn_cast<TensorType>(var_handle_op.getResource().getType());
    if (!tensor) return nullptr;
    TF::ResourceType resource =
        llvm::dyn_cast<TF::ResourceType>(tensor.getElementType());
    if (!resource || resource.getSubtypes().size() != 1) return nullptr;
    return resource.getSubtypes().front();
  } else if (auto global_tensor_op =

      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) {

        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)) {

      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

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) {