session_init_func.getLoc(), tensor_attr.getType(), tensor_attr);

  builder.create<TF::AssignVariableOp>(
      session_init_func.getLoc(), llvm::ArrayRef<mlir::Type>{},
      llvm::ArrayRef<mlir::Value>{var_handle_op_in_init->getResult(0),
                                  const_op.getResult()});
}

func::FuncOp CreateSessionInitFunc(ModuleOp module) {

                                 ShapedType weight_type) {
  ArrayRef<int64_t> input_shape = input_type.getShape();
  ArrayRef<int64_t> weight_shape = weight_type.getShape();

  const int64_t num_matmul_dim = 2;
  const int64_t num_input_batch_dim = input_type.getRank() - num_matmul_dim;
  const int64_t num_weight_batch_dim = weight_type.getRank() - num_matmul_dim;

  ArrayRef<int64_t> input_batch_dims =

  auto bad_consumers = b.getNamedAttr(
      "bad_consumers",
      b.getI32ArrayAttr(llvm::ArrayRef<int32_t>(
          versions.bad_consumers().data(),
          versions.bad_consumers().data() + versions.bad_consumers().size())));
  module->setAttr("tf.versions",
                  b.getDictionaryAttr(llvm::ArrayRef<mlir::NamedAttribute>(
                      {producer, min_consumer, bad_consumers})));
}

    state_vars.push_back(builder->create<TF::VarHandleOp>(
        loc,
        llvm::ArrayRef<Type>{RankedTensorType::get(
            {}, TF::ResourceType::get(llvm::ArrayRef<TensorType>{key_type},
                                      builder->getContext()))},
        llvm::ArrayRef<Value>{},
        llvm::ArrayRef<NamedAttribute>{
            builder->getNamedAttr(kDeviceAttr, builder->getStringAttr(device)),

}

// Transforms Conv2DBackPropFilter for space to depth.
void HandleConv2DBackPropFilter(TF::Conv2DBackpropFilterOp backprop,
                                ArrayRef<int32_t> old_filter_shape,
                                ArrayRef<int32_t> new_filter_shape,
                                int64_t block_size) {
  OpBuilder builder(backprop);
  builder.setInsertionPoint(backprop);

  auto arr = input_tensor.flat<T>();
  return ElementsAttr(mlir::DenseElementsAttr::get(
      type, llvm::ArrayRef(arr.data(), arr.size())));
}

ElementsAttr ConvertTensorOfCustomFloatType(const Tensor& tensor,
                                            RankedTensorType type) {
  auto buffer =
      llvm::ArrayRef(static_cast<char*>(tensor.data()), tensor.TotalBytes());
  return mlir::DenseElementsAttr::getFromRawBuffer(type, buffer);

  auto size = it->getSecond().size;
  OpBuilder builder(pop);
  auto new_buffer = builder.create<TF::IdentityOp>(
      pop.getLoc(), ArrayRef<Type>{buffer.getType()}, ArrayRef<Value>{buffer});
  auto new_size = builder.create<TF::SubOp>(
      pop.getLoc(), ArrayRef<Type>{size.getType()},
      ArrayRef<Value>{size, cutil::GetR1Const({1LL}, builder, pop.getLoc())});
  auto element = cutil::GetElement(new_size, new_buffer, builder, pop.getLoc());

  if (!x_ty || !y_ty) return success();

  ArrayRef<int64_t> x_shape = x_ty.getShape();
  ArrayRef<int64_t> y_shape = y_ty.getShape();

  llvm::SmallVector<int64_t, 4> result_batch_shape;
  llvm::ArrayRef<int64_t> x_batches = x_shape.drop_back(2);
  llvm::ArrayRef<int64_t> y_batches = y_shape.drop_back(2);

  if (auto loc = mlir::dyn_cast<NameLoc>(var_handle_op->getLoc()))
    return loc.getName().str();
  return var_handle_op.getSharedName().str();
}

std::vector<std::string> GetVariableNames(
    llvm::ArrayRef<TF::VarHandleOp> var_handle_ops) {
  std::vector<std::string> names;
  names.reserve(var_handle_ops.size());
  for (auto var_handle_op : var_handle_ops)
    names.push_back(GetVariableName(var_handle_op));
  return names;

IfRegionOp CreateOptionalDatasetIf(
    OpBuilder builder, ReduceDatasetOp reduce_dataset, FuncOp reduce_func,
    IteratorGetNextAsOptionalOp get_next, OptionalHasValueOp optional_has_value,
    ArrayRef<Value> body_args, ArrayRef<Type> dataset_types) {
  const Location loc = reduce_dataset.getLoc();
  // If returns are the state variables.
  SmallVector<Type, 4> if_return_types;
  const int state_size =