auto result_element_type_fn = [&](int index) {
    return mlir::cast<TensorType>(op->getResult(index).getType())
        .getElementType();
  };

  llvm::SmallVector<ShapedTypeComponents, 4> inferred_return_shapes;
  if (failed(InferReturnTypeComponentsForTFOp(
          /*location=*/std::nullopt, op, graph_version_, operand_as_constant_fn,
          op_result_as_shape_fn, result_element_type_fn,