return failure();
    }

    if (!IsI8F32UniformQuantizedPerAxisType(filter_type.getElementType())) {
      LLVM_DEBUG(llvm::dbgs() << "Expected a per-channel uniform quantized "
                                 "(i8->f32) type filter. Got: "
                              << filter_type.getElementType() << "\n");
      return failure();
    }

    if (!keys_ty || !keys_ty.hasStaticShape() ||
        !keys_ty.getElementType().isIntOrFloat())
      return rewriter.notifyMatchFailure(
          op,
          "only match for the case where the first operand has a static "
          "int/float shapeType");
    if (!indices_ty || !indices_ty.hasStaticShape() ||
        !indices_ty.getElementType().isInteger(32))
      return rewriter.notifyMatchFailure(
          op,

  const auto output_type = getType(0).cast<ShapedType>();

  // Folding only implemented for float tensors.
  if (!input_type.getElementType().isF32() ||
      !weights_type.getElementType().isF32() ||
      !output_type.getElementType().isF32() ||
      (has_bias && !bias_type.getElementType().isF32())) {
    return failure();
  }

  // Folding only implemented for static shapes

  if (!begin_attr || !size_attr) {
    return false;
  }

  auto begin_elem_ty = begin_attr.getType().getElementType();
  if (!begin_elem_ty.isInteger(32) && !begin_elem_ty.isInteger(64)) {
    return false;
  }
  auto size_elem_ty = size_attr.getType().getElementType();
  if (!size_elem_ty.isInteger(32) && !size_elem_ty.isInteger(64)) {
    return false;
  }

  // is up to ~2x faster.
  const bool is_f16 = input_ty.getElementType().isF16();
  if (is_f16 && CanUseTensorCores(devices)) return "NHWC";

  // For f32/f16 data type decision depends on the filter size in spatial
  // dimensions, for other data types we keep current data format.
  if (!input_ty.getElementType().isF32() && !input_ty.getElementType().isF16())
    return getDataFormat();

            size.getSExtValue(),
            start.getType().cast<TensorType>().getElementType()),
        start, limit, delta);
  }
  return RangeOp::build(
      builder, result,
      tensorflow::GetTypeFromTFTensorShape(
          {-1}, start.getType().cast<TensorType>().getElementType()),
      start, limit, delta);
}

OpFoldResult RangeOp::fold(FoldAdaptor adaptor) {

      mlir::dyn_cast<TensorFlowTypeWithSubtype>(operand_type.getElementType());
  if (!operand_handle_type) return result_type.getElementType();
  auto result_handle_type =
      mlir::cast<TensorFlowTypeWithSubtype>(result_type.getElementType());
  if (operand_handle_type.GetSubtypes().empty() ||
      !result_handle_type.GetSubtypes().empty())
    return result_type.getElementType();
  return operand_handle_type;
}

    } else {
      return std::nullopt;
    }
  }

  auto element_type = tensor_type.getElementType();
  tflite::TensorType tflite_element_type =
      GetTFLiteType(tensor_type.getElementType()).value();
  std::optional<std::vector<BufferOffset<tflite::VariantSubType>>>
      variant_params = BuildTFVariantType(element_type);

  DerivedTFLiteTypeAttr output_type = DerivedTFLiteTypeAttr<[{
    return getResult().getType().cast<TensorType>().getElementType().
        cast<IntegerType>().getWidth() > 32 ? tflite::TensorType_INT64 :
            tflite::TensorType_INT32;
    }], [{
      TypeAttr::get(getResult().getType().cast<TensorType>().getElementType())
    }]>;
}

def TFL_ArgMinOp : TFL_Op<"arg_min", [
    QuantizableResult,
    Pure]> {