"func.func"() ({
  ^bb0(%arg0: tensor<i32>):
    "func.return"(%arg0) : (tensor<i32>) -> ()
  }) {sym_name = "main_while_body_4225150", sym_visibility = "private", tf._input_shapes = [#tf_type.shape<>], tf.signature.is_stateful, function_type = (tensor<i32>) -> (tensor<i32>)} : () -> ()

              .getOutputDimensions()
              .getDefiningOp<shape::ShapeOfOp>();
      // Check if the shape come from the original conv op.
      if (!shape_of_op ||
          shape_of_op.getArg().getDefiningOp<mhlo::ConvolutionOp>() !=
              conv_op) {
        return failure();
      }
      Value new_shape_of = rewriter.create<shape::ShapeOfOp>(
          mul_op.getLoc(), shape_of_op.getType(), new_conv);

//  xla::Shape shape(xla::F32, {2, 2}, {}, {})
//  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =
//      tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(stream.parent());
//  done_event->Init();
//  Tensor dest_cpu_tensor;
//
//  XlaHostRecvDeviceContext device_context(&stream, gpu_dst,
//    shape, done_event);
//  device_context.CopyDeviceTensorToCPUSync(

                        llvm::SmallVectorImpl<int64_t>* shape) {
  shape->reserve(input_shape.dims());
  for (const auto& d : input_shape) {
    shape->push_back(d.size == kTFDynamicSize ? ShapedType::kDynamic : d.size);
  }
}

Status ConvertToMlirShape(const TensorShapeProto& input_shape,
                          llvm::SmallVectorImpl<int64_t>* shape) {
  shape->reserve(input_shape.dim_size());

}

// Helper method that given 'shape' and 'current_index' representing
// index in broadcasted tensor, get the index in the flat original tensor.
// 'shape' is computed from the original shape and the broadcast dimensions to
// match result shape.
int64_t GetElementIndex(llvm::SmallVectorImpl<int64_t> &shape,
                        llvm::SmallVectorImpl<int64_t> &current_index) {
  int64_t ind = 0;

  TF_RETURN_IF_ERROR(varhandle_op->SetAttrType("dtype", dtype));

  // Note that if shape is unknown rank, shape.dim_sizes() will be empty, and
  // shape.dims() will be -1.
  absl::InlinedVector<int64_t, 4UL> dim_sizes = shape.dim_sizes();
  TF_RETURN_IF_ERROR(varhandle_op->SetAttrShape(
      "shape", reinterpret_cast<const int64_t*>(dim_sizes.data()),
      shape.dims()));
  TF_RETURN_IF_ERROR(varhandle_op->SetAttrString("container", "", 0));

// 'shape' is the original shape with padding to match result shape.
int64_t GetElementIndex(const std::vector<int64_t>& shape,
                        const std::vector<int64_t>& current_index) {
  int64_t ind = 0;
  int64_t mul = 1;
  for (int i = shape.size() - 1; i >= 0; --i) {
    ind += (current_index[i] % shape[i]) * mul;
    mul *= shape[i];
  }
  return ind;
}

node {
  name: "Placeholder"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 2
        }
        dim {
          size: 5
        }
        dim {
          size: 3
        }
      }
    }
  }
}
node {

  // tail of the other operand and the intermediate result isn't used by other
  // ops.
  // $rhs is required to be the tail shape of $lhs, so after transformation the
  // shape of the binary op result is valid. For example, assume the shapes of
  // $input, $lhs and $rhs are [1600], [1,40,40] and [40x1]. After the
  // transformation, the shape of the binary op result is [40x1600], which

// Fill in the contents of TensorShapeProto for the given shape.
// ShapeContainerT is any type with the following methods:
//   bool hasRank()
//   ArrayRef<int64_t> getShape()
// This includes mlir::TF::ShapeAttr and mlir::ShapedType.
template <typename ShapeContainerT>
void SetTensorShapeProto(ShapeContainerT shape, TensorShapeProto* proto) {
  if (shape.hasRank()) {
    for (int64_t dim : shape.getShape()) {