#include "mlir/IR/Operation.h"  // from @llvm-project
#include "tensorflow/compiler/tf2xla/tf2xla_defs.h"

namespace mlir {
namespace TF {

using ::tensorflow::kValidDeviceTypes;

LogicalResult HasValidCompilationAndReplicationAttributes(Operation& op) {
  auto replicate_attr = op.getAttrOfType<StringAttr>(kReplicationInfoAttr);
  auto compile_attr = op.getAttrOfType<StringAttr>(kCompileDeviceTypeAttr);

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_structs.h"

      Value stb_block_shape, Value bts_block_shape, int64_t expand_axis,
      PatternRewriter& rewriter) const;

 public:
  LogicalResult matchAndRewrite(Conv2dOpTy op,
                                PatternRewriter& rewriter) const override;
};

template <typename Conv2dOpTy>
LogicalResult ConvertTFDilatedConvOp<Conv2dOpTy>::matchAndRewrite(
    Conv2dOpTy op, PatternRewriter& rewriter) const {
  if (!op.getResult().hasOneUse()) {

  });

  return failure(walk_result.wasInterrupted());
}

}  // anonymous namespace

namespace TF {

LogicalResult CleanupAndCanonicalizeForResourceOpLifting(func::FuncOp func) {
  return CleanupAndCanonicalize(func);
}

LogicalResult CleanupAndCanonicalizeForResourceOpLifting(ModuleOp module) {
  auto walk_result = module.walk([](tf_device::ClusterOp cluster) {

  // TF op.
  LogicalResult CollectInputsAndAttributes(
      PatternRewriter& rewriter, TFRFuncOp signature, CallOp call_op,
      SmallVectorImpl<Value>* inputs, NamedAttrList* arg_attrs,
      llvm::StringMap<Attribute>* derived_attrs) const;

  // Uses the collected attribute values to derive all the output types.
  LogicalResult DeriveOutputTypes(Location loc, FunctionType signature,

};

LogicalResult DecomposeStackOpsInternal(
    Block*, ModuleOp, llvm::SmallDenseMap<Value, Value>*,
    llvm::StringMap<PartitionedCallStackOpsInfo>*);

// Handles stack usage by a tf.While. It will convert the body and conditional
// function signatures, and performs stack ops decomposition on them.
LogicalResult HandleWhileOp(
    TF::WhileOp while_op, ModuleOp module,

#include "mlir/IR/Location.h"  // from @llvm-project
#include "mlir/IR/Types.h"  // from @llvm-project
#include "mlir/IR/Value.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/client/sharding_builder.h"
#include "xla/service/hlo_parser.h"
#include "xla/xla_data.pb.h"

namespace tensorflow {

// is the associated virtual device name (`TPU_REPLICATED_HOST_i`).
mlir::LogicalResult GetDeviceToHostMap(
    mlir::tf_device::ClusterOp cluster,
    llvm::SmallVector<std::string, 8>& core_to_host);

// Returns the first TPU device, for use in the non-replicated case. The list of
// TPU devices is retrived from `op`'s module ancestor.
mlir::LogicalResult GetNonReplicatedTPU0(mlir::Operation* op,

#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

#include "mlir/Interfaces/CallInterfaces.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/utils/lstm_utils.h"