Status DefaultPaddedShapeFn(const Tensor& tensor, xla::Shape* shape) {
  const tensorflow::XlaTensor* xla_tensor =
      tensorflow::XlaTensor::FromTensor(&tensor);
  if (xla_tensor == nullptr) {
    return TensorShapeToXLAShape(tensor.dtype(), tensor.shape(), shape);
  }

  const xla::ShapedBuffer& shaped_buffer = xla_tensor->shaped_buffer();
  *shape = shaped_buffer.on_device_shape();
  return absl::OkStatus();
}

    "'h_scale' type: 'float'} attr : { name: 'max_classes_per_detection' "
    "type: 'int'} attr : { name: 'max_detections' type: 'int'} attr : { "
    "name: 'nms_iou_threshold' type: 'float'} attr : { name: "
    "'nms_score_threshold' type: 'float'} attr : { name: 'num_classes' type: "
    "'int'} attr : { name: 'w_scale' type: 'float'} attr : { name: 'x_scale' "

  CreateContext();

  std::vector<XlaCompiler::Argument> args(2);
  args[0].kind = XlaCompiler::Argument::kParameter;
  args[0].type = DT_INT32;
  args[0].shape = TensorShape({1, 3});
  args[1].kind = XlaCompiler::Argument::kParameter;
  args[1].type = DT_INT32;
  args[1].shape = TensorShape({1, 3});

  const XlaCompiler::CompilationResult* result;
  xla::PjRtLoadedExecutable* executable;

         << compilation_result.xla_output_shape.ToString() << '\n';

  for (const auto& xla_output_description : compilation_result.outputs) {
    output << "// XlaOutputDescription type="
           << DataTypeString(xla_output_description.type) << " shape=("
           << absl::StrJoin(xla_output_description.shape.dim_sizes(), ", ")
           << ')';
    if (xla_output_description.input_index >= 0)

      // Args
      {"x: T"},
      // Return values
      {"y: T"},
      // Attr def
      {"T: {float, double, int32, int64}"},
      // Nodes
      {
          {{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
          {{"scale"}, "Cast", {"two"}, {{"SrcT", DT_INT64}, {"DstT", "$T"}}},
          {{"y"}, "Mul", {"x", "scale"}, {{"T", "$T"}}},
      });
}

class MlirGraphOptimizationPassTest : public Test {

  };
  TF_EXPECT_OK(ShapeAnnotationsMatch(graph, shape_info, expected));
}

TEST(ShapeInferenceTest, WhileLoopWithResource) {
  // Graph:
  // x = resource_variable_ops.var_handle_op(dtype=dtypes.float32, shape=[2, 3])
  // y = control_flow_ops.while_loop(lambda _: true, lambda x: x, [x])
  Graph graph(OpRegistry::Global());
  {
    Scope scope = Scope::NewRootScope().ExitOnError();

    auto x =

  // When set, quantized_bias_type defines the dtype for both bias and accumulator.
  quantized_bias_type: TensorType;
}

table ExpandDimsOptions {
}

table SparseToDenseOptions {
  validate_indices:bool;
}

table EqualOptions {
}

table NotEqualOptions {
}

table ShapeOptions {
  // Optional output type of the operation (int32 or int64). Defaults to int32.

                         kSavedModelVariablesFilename);
        ImmediateTensorHandlePtr restored_output;
        TF_RETURN_IF_ERROR(internal::SingleRestore(
            context, variables_path_prefix, checkpoint_key, variable->dtype(),
            &restored_output));

        // Assign the restored tensor's value to the variable
        return variable->Assign(restored_output.get());
      }));

  return Status();
}

    auto iter = placeholders.find(placeholder_index);
    Node* placeholder_node;
    if (iter == placeholders.end()) {
      NodeDefBuilder placeholder_builder(new_name, "Placeholder");
      placeholder_builder.Attr("dtype", src->output_type(src_output));
      string outside_compilation_attr;
      TF_RETURN_IF_ERROR(GetNodeAttr(dst->attrs(),
                                     outside_compilation_attr_name,

}

// -----

// Tests that the pass can transform execution with model parallelism and no
// replication.
//
// The following TPUCompileMetadataProto is used:
// args {
//   dtype: DT_FLOAT
//   shape {
//     dim {
//       size: 128
//     }
//   }
// }
// num_replicas: 1
// num_cores_per_replica: 2

// CHECK-LABEL: func @parallel_execute