"Shape must be fully-defined, got: shape[", dim, "] = ", shape[dim]));
    }
    num_elements *= shape[dim];
  }
  if (data.size() != num_elements) {
    return tensorflow::errors::InvalidArgument(absl::StrCat(
        "Mismatched shape and data size: \n", "Shape num_elements: ",
        num_elements, "\n", "Data size: ", data.size(), "\n"));
  }

// Extracts shape from XlaArgument as TensorShape. If shape is a xla::Shape,
// that is converted to a TensorShape.
absl::StatusOr<TensorShape> GetTensorShapeFromXlaArgument(
    const XlaArgument& arg) {
  if (absl::holds_alternative<xla::Shape>(arg.shape)) {
    TensorShape arg_shape;
    TF_RETURN_IF_ERROR(
        XLAShapeToTensorShape(std::get<xla::Shape>(arg.shape), &arg_shape));
    return arg_shape;

  if (input_type_size < output_type_size) {
    TF_ShapeInferenceContextWithRankAtLeast(ctx, shape, 1, shape, status);

    if (TF_GetCode(status) == TF_OK) {
      TF_DimensionHandle* last_dim = TF_NewDimensionHandle();
      size_t divisor_val = output_type_size / input_type_size;
      TF_ShapeInferenceContextDim(ctx, shape, -1, last_dim);
      if (!TF_DimensionHandleValueKnown(last_dim) ||

  dltensor_in->device = {kDLCPU, 0};
  dltensor_in->ndim = static_cast<int32_t>(shape.size());
  dltensor_in->dtype = {kDLFloat, 32, 1};
  dltensor_in->shape = shape.data();
  dltensor_in->strides = strides.data();
  TFE_TensorHandle* handle = TFE_HandleFromDLPack(&dlm_in, status, ctx);
  ASSERT_NE(handle, nullptr)
      << TF_Message(status) << " (shape=[" << absl::StrJoin(shape, ",")
      << "], strides=[" << absl::StrJoin(strides, ",") << "])";

    {
      "tensors": [
        {
          "shape": [1, 5, 2],
          "name": "input0"
        },
        {
          "shape": [2, 5],
          "buffer": 1,
          "name": "input2input_weights1"
        },
        {
          "shape": [2, 5],
          "buffer": 2,
          "name": "input2forget_weights2"
        },
        {
          "shape": [2, 5],
          "buffer": 3,

    // Merge node causes a loop so we remove NextIteration->Merge edge before
    // performing shape inference. But removing those edges also prevents us
    // from inferring output shape for Merge node (we need shapes for all its
    // inputs).
    // For loop invariant resource input's Merge node, we set output resource
    // shape as Enter node's resource shape.
    // TODO(b/129367850): clean this up.

// Shape[i].Score represents function values at meeting points.
//
// function f(p) is defined as:
//
//	shape[0].Score for p < shape[0].Utilization
//	shape[n-1].Score for p > shape[n-1].Utilization
//
// and linear between points (p < shape[i].Utilization)
func BuildBrokenLinearFunction(shape FunctionShape) func(int64) int64 {
	return func(p int64) int64 {
		for i := 0; i < len(shape); i++ {
			if p <= int64(shape[i].Utilization) {

  @tf.function(input_signature=[])
  def f0001_multiple_results_no_punctuation(self):
    return tf.constant(1.0, shape=[1]), tf.constant(1.0, shape=[2])

  # Check index paths for results written explicitly with parentheses.
  # This is semantically equivalent to the earlier test without parentheses,

//  se::DeviceMemoryBase gpu_dst{device_tensor.data(), 4 * sizeof(float)};
//  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();
//
//  XlaHostSendDeviceContext device_context(&stream, &gpu_dst,
//    shape, done_event);
//  device_context.CopyCPUTensorToDeviceSync(

// CHECK-NEXT: %[[SHAPE:.*]] = "tf.RiscShape"(%arg0) {T = i32} : (tensor<2x3xf32>) -> tensor<*xi32>
// CHECK-NEXT: %[[ALPHA1:.*]] = "tf.RiscBroadcast"(%[[ALPHA]], %[[SHAPE]]) : (tensor<f32>, tensor<*xi32>) -> tensor<*xf32>
// CHECK-NEXT: %[[MAX:.*]] = "tf.RiscMaximum"(%arg0, %[[ALPHA1]]) : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<*xf32>