}

    rewriter.replaceOp(op, results);
    return success();
  }
};

// Canonicalize ShapeNOp to ShapeOp if there is only one operand.
class ShapeNToShape : public OpRewritePattern<ShapeNOp> {
  using OpRewritePattern<ShapeNOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(ShapeNOp op,
                                PatternRewriter &rewriter) const override {
    if (op.getNumOperands() != 1) {

  auto element_ty = input_ty.getElementType();
  output_ty = UnrankedTensorType::get(element_ty);

  auto shape_ty = shape.getType().dyn_cast<RankedTensorType>();
  if (!shape_ty) return success();
  if (shape_ty.getRank() != 1)
    return error_handler(llvm::formatv(
        "requires 'shape' to be rank 1, but got {0}", shape_ty.getRank()));

  DenseIntElementsAttr shape_attr;
  if (!matchPattern(shape, m_Constant(&shape_attr))) {

    llvm::ArrayRef<int64_t> shape_ref =
        mlir::cast<TensorType>(tensor_attr.getType()).getShape();
    if (mlir::failed(check_shape(shape_ref))) return std::nullopt;

    shape = std::vector<int32_t>(shape_ref.begin(), shape_ref.end());
  } else if (type.hasRank()) {
    llvm::ArrayRef<int64_t> shape_ref = type.getShape();
    if (mlir::failed(check_shape(shape_ref))) return std::nullopt;

  // CHECK-LABEL:select_v2_with_high_dims_dynamic_shape_both_sides
  // CHECK: %[[SHAPE_0:.*]] = "tfl.shape"(%arg0) : (tensor<8x7x6x5x?x3x2x1xf32>) -> tensor<8xi64>
  // CHECK: %[[SHAPE_1:.*]] = "tfl.shape"(%arg1) : (tensor<?x3x2x1xf32>) -> tensor<4xi64>
  // CHECK: %[[BROADCAST_ARGS_0:.*]] = "tfl.broadcast_args"(%[[SHAPE_0]], %[[SHAPE_1]]) : (tensor<8xi64>, tensor<4xi64>) -> tensor<8xi64>

      mhlo::GetDimensionSizeOp op, OpAdaptor adaptor,
      ConversionPatternRewriter& rewriter) const final {
    ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
    Value shape_op = rewriter.create<TF::ShapeOp>(op.getLoc(), op.getOperand(),
                                                  rewriter.getBoolAttr(true));
    Value size =
        BuildIntArrayConstOp(builder, rewriter, llvm::SmallVector<int64_t>({1}),

    %tl_1 = "tf.TensorListPushBack"(%tl_0, %elem) : (tensor<!tf_type.variant<tensor<?x1xf32>>>, tensor<16x1xf32>) -> tensor<!tf_type.variant<tensor<?x1xf32>>>
    %shape_32 = "tf.TensorListElementShape"(%tl_1) : (tensor<!tf_type.variant<tensor<?x1xf32>>>) -> tensor<?xi32>
    %shape_64 = "tf.TensorListElementShape"(%tl_1) : (tensor<!tf_type.variant<tensor<?x1xf32>>>) -> tensor<?xi64>

      return ComputeOutputComponent(ValuePort(id.getInput()), values);
    return nullptr;
  }

  if (auto shape_op = dyn_cast<TF::ShapeOp>(op)) {
    // No shape available in an unranked tensor type.
    auto operand_ty =
        mlir::dyn_cast<RankedTensorType>(shape_op.getOperand().getType());
    if (!operand_ty) return nullptr;

    // Shape op has a single output so the first element should always be zero

      shape.push_back(dim.getSExtValue());
    }
    return tensorflow::GetTypeFromTFTensorShape(shape, etype);
  }

  if (auto shape_op = dims.getDefiningOp<ShapeOp>()) {
    if (auto t = mlir::dyn_cast<ShapedType>(shape_op.getInput().getType())) {
      return t;
    }
  }

  return UnrankedTensorType::get(etype);
}

absl::StatusOr<mlir::Type> ImporterBase::InferOutputType(
    const Node& node, int idx, mlir::Builder builder) {
  DataType dtype = node.properties()->output_types[idx];

  // Returns output type given inference context.
  auto shape_ic =
      [&](shape_inference::InferenceContext* c) -> absl::StatusOr<mlir::Type> {
    // TODO(b/200093974): Post triage, consider following
    // GraphConstructor::ValidateShape in checking _output_shapes always.