if (pack_axis < 0) {
    pack_axis += rank;
  }

  // Concat out shape.
  for (int i = 0; i < rank; ++i) {
    int64_t dim_size = input_type.getDimSize(i);
    if (i == pack_axis) {
      dim_size *= count;
    }
    concat_out_shape.push_back(dim_size);
  }

  // Pack out shape.
  int j = 0;
  for (int i = 0; i < rank + 1; ++i) {
    if (i == pack_axis) {

  }
  func.return
}

// CHECK: %[[SHAPE:[0-9]*]] = "tf.Shape"(%[[ARG_0]])
// CHECK: %[[OP_A:[0-9]*]] = "tf.opA"(%[[SHAPE]])
// CHECK: %[[OP_B:[0-9]*]] = "tf.opB"(%[[SHAPE]], %[[OP_A]])
// CHECK: tf_device.replicate
// CHECK:   tf_device.return %[[SHAPE]], %[[OP_A]], %[[OP_B]]


// CHECK-LABEL: func @nested_ops

		// if either host or share name is different, return s
		{`\\host\share\foo`, ``, `\\host\share\foo`},
		{`\\host\share\foo`, `\foo`, `\\host\share\foo`},
		{`\\host\share\foo`, `\\host\other\`, `\\host\share\foo`},
		{`\\host\share\foo`, `\\other\share\`, `\\host\share\foo`},
		{`\\host\share\foo`, `\\host\`, `\\host\share\foo`},
		{`\\host\share\foo`, `\share\`, `\\host\share\foo`},

		// only volume names are case-insensitive

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

  return llvm::to_vector(llvm::map_range(shapes, [](int64_t shape) {
    return shape == kTFDynamicSize ? mlir::ShapedType::kDynamic : shape;
  }));
}

llvm::SmallVector<int64_t> ConvertMlirShapeToTF(
    llvm::ArrayRef<int64_t> shapes) {
  return llvm::to_vector(llvm::map_range(shapes, [](int64_t shape) {
    return mlir::ShapedType::isDynamic(shape) ? kTFDynamicSize : shape;
  }));
}

   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,