absl::Span<AbstractTensorHandle*> outputs) {
  PartialTensorShape shape;
  TF_RETURN_IF_ERROR(inputs[0]->Shape(&shape));
  if (shape.dims() != 0) {
    return errors::InvalidArgument(
        "Tensor expected to have scalar shape found rank: ", shape.dims());
  }
  return absl::OkStatus();
}

TEST_P(UnifiedAPI, TestTensorShapeScalar) {
  if (UseFunction() && UseMlir()) {

  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
  auto y = AddN(scope_, xs);
  RunTest(xs, {shape, shape, shape}, {y}, {shape});
}

TEST_F(NaryGradTest, Add) {
  TensorShape x1_shape({3, 2, 5});
  TensorShape x2_shape({2, 5});
  auto x1 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x1_shape));
  auto x2 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x2_shape));

    TF_RETURN_IF_ERROR(session->Create(graph));
    RunOptions run_options;
    Tensor input_a = CreateInputTensor(shape, 0);
    Tensor input_b = CreateInputTensor(shape, shape.num_elements());
    Tensor input_c = CreateInputTensor(shape, 2 * shape.num_elements());
    TF_RETURN_IF_ERROR(session->Run(
        run_options,
        {std::make_pair("a", input_a), std::make_pair("b", input_b),

   node [fontname = "Sans"];
   edge [fontname = "Sans"];

   "apiElements" [shape="box", label=<<B>apiElements</B>(C)>, color="pink", bgcolor="pink", style="filled"]
   "apiElements" -> "api"
   "apiElements" -> "compileOnlyApi"

   "api" [shape="box", label=<<B>api</B>>, color="chartreuse3", bgcolor="chartreuse3", style="filled"]
   "compileOnlyApi" [shape="box", label=<<B>compileOnlyApi</B>>, color="chartreuse3", bgcolor="chartreuse3", style="filled"]

// CHECK-NEXT:     }, {
// CHECK-NEXT:       shape: [ 0 ],
// CHECK-NEXT:       name: "input_to_input_intermediate",
// CHECK-NEXT:       has_rank: true
// CHECK-NEXT:     }, {
// CHECK-NEXT:       shape: [ 0 ],
// CHECK-NEXT:       name: "input_to_forget_intermediate",
// CHECK-NEXT:       has_rank: true
// CHECK-NEXT:     }, {
// CHECK-NEXT:       shape: [ 0 ],
// CHECK-NEXT:       name: "input_to_cell_intermediate",

  auto c_api_client = down_cast<xla::PjRtCApiClient*>(pjrt_client->get());
  std::vector<int32_t> data(1, 0);
  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::S32, {1});

  auto buffer = c_api_client->pjrt_c_client()->client->BufferFromHostBuffer(
      data.data(), shape.element_type(), shape.dimensions(),
      /*byte_strides=*/std::nullopt,
      xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall, nullptr,

// CHECK-NEXT:   subgraphs: [ {
// CHECK-NEXT:     tensors: [ {
// CHECK-NEXT:       shape: [ 1, 4 ],
// CHECK-NEXT:       buffer: 1,
// CHECK-NEXT:       name: "arg0",
// CHECK-NEXT:       quantization: {
// CHECK-EMPTY:
// CHECK-NEXT:       },
// CHECK-NEXT:       has_rank: true
// CHECK-NEXT:     }, {
// CHECK-NEXT:       shape: [ 4, 4 ],
// CHECK-NEXT:       buffer: 2,
// CHECK-NEXT:       name: "arg1",

// CHECK-NEXT:   subgraphs: [ {
// CHECK-NEXT:     tensors: [ {
// CHECK-NEXT:       shape: [ 1, 4 ],
// CHECK-NEXT:       buffer: 1,
// CHECK-NEXT:       name: "arg0",
// CHECK-NEXT:       quantization: {
// CHECK-EMPTY:
// CHECK-NEXT:       },
// CHECK-NEXT:       has_rank: true
// CHECK-NEXT:     }, {
// CHECK-NEXT:       shape: [ 4, 4 ],
// CHECK-NEXT:       buffer: 2,
// CHECK-NEXT:       name: "arg1",

  return transposed_type;
}

// Returns shape of a ranked tensor.
// Precondition: output_val's is ranked tensor.
// Returns a truncated shape when `truncate` is set to true.
inline DenseElementsAttr GetShape(Value output_val, bool truncate = false) {
  auto output_shape = output_val.getType().dyn_cast<ShapedType>().getShape();

  SmallVector<int32_t> shape;
  shape.reserve(output_shape.size());

import gradlebuild.basics.decapitalize
import gradlebuild.shade.ArtifactTypes.buildReceiptType
import gradlebuild.shade.ArtifactTypes.classTreesType
import gradlebuild.shade.ArtifactTypes.entryPointsType
import gradlebuild.shade.ArtifactTypes.manifestsType
import gradlebuild.shade.ArtifactTypes.relocatedClassesAndAnalysisType
import gradlebuild.shade.ArtifactTypes.relocatedClassesType
import gradlebuild.shade.extension.ShadedJarExtension