CompileTimeConstantOperand<TF::TensorArrayV3Op, 0>,       // $size
      // $element_shape
      CompileTimeConstantOperand<TF::TensorListFromTensorOp, 1>,
      // $element_shape, $num_elements
      CompileTimeConstantOperand<TF::TensorListReserveOp, 0, 1>,
      // $begin, $end, $strides
      CompileTimeConstantOperand<TF::TensorStridedSliceUpdateOp, 1, 2, 3>,

    if (src->name() != dst->name()) {
      Status s = CheckIfTPUInitialized();
      if (!s.ok()) {
        done(s);
        return absl::OkStatus();
      }
    }
    if (input->shape().num_elements() == 0) {
      // Zero-element tensors have no backing buffers.
      done(absl::OkStatus());
      return absl::OkStatus();
    }

    se::Stream* const src_compute_stream = src_xla_context->stream();

  // We can create an MLIR Attribute more efficiently in this case.
  TensorShape input_tensor_shape(input_tensor.tensor_shape());
  if (NumberOfMaterializedElements(input_tensor) == 1 &&
      input_tensor_shape.num_elements() > 1) {
    // We first convert this TensorProto to one of shape [1]. We then create an
    // Attribute for that proto, and finally splat the Attribute.

    TensorProto tensor_copy = input_tensor;

      %outputs_55, %control_56 = tf_executor.island wraps "tf.TensorListStack"(%outputs_53#3, %outputs_27) {_tpu_replicate = "cluster_sample_sequence", device = "", num_elements = 1024 : i64} : (tensor<!tf_type.variant>, tensor<1xi32>) -> tensor<*xf32>

TypedAttr GetNumElementsOrOne(Type type) {
  auto shaped_type = mlir::cast<ShapedType>(type);
  int32_t num_elements =
      shaped_type.hasStaticShape() ? shaped_type.getNumElements() : 1;

  OpBuilder builder(type.getContext());

  return DenseIntElementsAttr::get(
      RankedTensorType::get({}, builder.getI32Type()),
      {llvm::APInt(32, num_elements, true)});
}

// Reshapes value to a given shape.

                                    builder->getBoolAttr(false));
}

// Creates a mhlo.rng_uniform op with `builder` to generate `num_elements`
// 32-bit integer numbers in the range of [`lower_limit`, `upper_limit`).
static mhlo::RngOp CreateRngUniform32(Location loc, int num_elements,
                                      int lower_limit, int upper_limit,
                                      OpBuilder *builder) {

  const int num_dims = 2;
  int64_t* dims = new int64_t[num_dims];
  int64_t num_elements = 1;
  dims[0] = batch_size;
  dims[1] = 1;
  for (int64_t i = 0; i < num_dims; ++i) {
    num_elements *= dims[i];
  }
  TF_Tensor* t = TF_AllocateTensor(TF_STRING, dims, num_dims,
                                   sizeof(TF_TString) * num_elements);
  delete[] dims;

// Construct the tensor for the given type and buffer.
static Tensor MakeTensor(DataType dtype, const TensorShape& shape,
                         se::DeviceMemoryBase buffer, Allocator* allocator) {
  size_t expected_size = shape.num_elements() * DataTypeSize(dtype);
  auto* tensor_buffer = new XlaTensorBuffer(buffer.opaque(), expected_size,
                                            buffer.size(), allocator);
  Tensor t(dtype, shape, tensor_buffer);

  let description = [{
element_shape: the shape of the future elements of the list
num_elements: the number of elements to reserve
handle: the output list
element_dtype: the desired type of elements in the list.
  }];

  let arguments = (ins
    TF_I32OrI64Tensor:$element_shape,
    TF_Int32Tensor:$num_elements
  );
}

  %cst = arith.constant dense<3> : tensor<1xi32>
  %0 = "tf.TensorListFromTensor"(%arg0, %cst) : (tensor<?x3xf32>, tensor<1xi32>) -> tensor<!tf_type.variant<tensor<3xf32>>>
  %1 = "tf.TensorListStack"(%0, %cst) {num_elements = 2 : i64} : (tensor<!tf_type.variant<tensor<3xf32>>>, tensor<1xi32>) -> tensor<2x3xf32>
  func.return %1 : tensor<2x3xf32>

// CHECK:  [[ELEM_SHAPE:%cst.*]] = arith.constant dense<3> : tensor<1xi32>