func.return
}

// -----

func.func @testConcatOffest(%concat_dim: tensor<i32>, %shape0: tensor<3xi32>, %shape1: tensor<3xi32>) {
  // expected-error @+1 {{'tf.ConcatOffset' op requires operand and result 1 to have compatible shapes}}

  //
  // For example,
  //
  //   %shape  = tf.Shape(%arg)                   // %arg: tensor<?x2x3x1xf32>
  //   %height = tf.StridedSlice(%shape, 1, 2, 1)
  //
  // In this case %height can be replaced with a constant 2.
  //
  // Or,
  //
  //   %shape  = tf.Shape(%arg)                   // %arg: tensor<?x2x3x1xf32>
  //   %spatial_shape = tf.StridedSlice(%shape, 1, 3, 1)
  //

  {
    GraphDefBuilder shape2(GraphDefBuilder::kFailImmediately);
    Node* key_constant = KeyPlaceholder("F1", shape2.opts());
    Node* recv1 = RecvAtHost(
        ops::NodeOut(key_constant, 0), "F1", "F1", "O1", {DT_FLOAT, DT_FLOAT},
        shape2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
    Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1),
                     shape2.opts()

  }

  std::vector<std::optional<std::vector<int>>> shapes;
  TF_RETURN_IF_ERROR(::tensorflow::ParseNodeShapes(input_shapes, shapes));

  for (const auto& shape : shapes) {
    if (!shape) {
      return absl::AbortedError("Missing input argument shapes");
    }
    parsed_shapes.push_back(SmallVector<int64_t>(shape->begin(), shape->end()));
  }
  return parsed_shapes;
}

// 'shape' is the original shape with padding to match result shape.
int64_t GetElementIndex(const std::vector<int64_t>& shape,
                        const std::vector<int64_t>& current_index) {
  int64_t ind = 0;
  int64_t mul = 1;
  for (int i = shape.size() - 1; i >= 0; --i) {
    ind += (current_index[i] % shape[i]) * mul;
    mul *= shape[i];
  }
  return ind;
}

  // Check all elements besides at concat_dim match across all shape tensors.
  SmallVector<int32_t, 4> shape0;
  shape0.reserve(num_dims);
  for (int32_t dim : shapes.front().getValues<int32_t>()) shape0.push_back(dim);

  for (DenseIntElementsAttr shape : llvm::drop_begin(shapes, 1)) {
    for (const auto& dims_and_idx : llvm::enumerate(llvm::zip(shape0, shape))) {
      if (dims_and_idx.index() == concat_dim) continue;

  std::vector<PartialTensorShape> shapes;
  shapes.reserve(num_shapes);
  for (int i = 0; i < num_shapes; ++i) {
    if (num_dims[i] < 0) {
      shapes.emplace_back();
    } else {
      shapes.emplace_back(ArraySlice<int64_t>(
          reinterpret_cast<const int64_t*>(dims[i]), num_dims[i]));
    }
  }
  desc->node_builder.Attr(attr_name, shapes);
}

  let options = [
    Option<"max_iterations_", "max-iterations", "int64_t", /*default=*/"10",
           "Maximum shape inference iterations">,
    Option<"input_arg_shapes_", "input-arg-shapes", "std::string", /*default=*/"",
           "Input tensor shapes. Shapes for different tensors are separated by ':', and dimension sizes for the same tensor are separated by ','">,
  ];
}

    }

    std::vector<PartialTensorShape> shapes;
    if (!GetNodeAttr(e->src()->attrs(), kXlaInferredShapesAttrName, &shapes)
             .ok()) {
      return std::nullopt;
    }

    const PartialTensorShape shape = shapes[e->src_output()];
    if (!shape.IsFullyDefined()) {
      return std::nullopt;
    }

    results[e->dst_input()] = shape;
  }
  return results;
}

}

// This assumes that the bias is of shape NxCx1x1 and doesn't require transpose
// Its corresponding constraint is optimize_patterns.td:IsBiasShape()
ElementsAttr ReshapeNCHWBiasToNHWC(Value v, Attribute a) {
  auto elements = mlir::cast<DenseElementsAttr>(a);
  auto shape = mlir::cast<ShapedType>(v.getType()).getShape();
  if (shape.size() != 4 || shape[2] != 1 || shape[3] != 1) return elements;