Type shape_dtype = getElementTypeOrSelf(op.getElementShape().getType());
    Value num_elements = operands[1];
    IntegerAttr attr;
    if (matchPattern(num_elements, m_Constant(&attr))) {
      return CreateI32SplatConst(op.getLoc(), rewriter, {1}, attr.getInt());
    }
    if (auto const_op = num_elements.getDefiningOp<TF::ConstOp>()) {
      return CreateI32SplatConst(op->getLoc(), rewriter, {1},

  // CHECK-NEXT: TensorListStack
  %2 = "tf.TensorListStack"(%1#2, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 2 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<2x8xf32>
  %3 = "tf.TensorListStack"(%1#3, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 2 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<2x8xf32>

  SmallVector<ValueType, 16> new_values;
  new_values.reserve(num_elements);
  ValueType new_value = start;
  for (int i = 0; i < num_elements; ++i) {
    new_values.push_back(new_value);
    new_value = new_value + delta;
  }
  // Result is always a 1-D tensor.
  auto new_result_type =
      tensorflow::GetTypeFromTFTensorShape({num_elements}, result_elem_type);
  return DenseElementsAttr::get(new_result_type, new_values);

  if (h == nullptr) {
    status->status = tensorflow::errors::InvalidArgument("Invalid handle");
    return -1;
  }

  int64_t num_elements = -1;
  status->status = tensorflow::unwrap(h)->NumElements(&num_elements);
  return num_elements;
}

int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, int dim_index,
                            TF_Status* status) {
  if (h == nullptr) {

    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;

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);