PR 3004 Suggested Fixes As per PR 3000 updating missing #include statements (IWYU)

examples/cpp/distributed.cpp

@@ -1,8 +1,13 @@
 // Copyright © 2024 Apple Inc.
 
 #include <iostream>
+#include <optional>
 
-#include "mlx/mlx.h"
+#include "mlx/array.h"
+#include "mlx/distributed/distributed.h"
+#include "mlx/distributed/ops.h"
+#include "mlx/ops.h"
+#include "mlx/utils.h"
 
 namespace mx = mlx::core;
 

examples/cpp/linear_regression.cpp

@@ -2,9 +2,15 @@
 
 #include <chrono>
 #include <cmath>
+#include <functional>
 #include <iostream>
+#include <ratio>
 
-#include "mlx/mlx.h"
+#include "mlx/array.h"
+#include "mlx/ops.h"
+#include "mlx/random.h"
+#include "mlx/transforms.h"
+#include "mlx/utils.h"
 #include "timer.h"
 
 /**

examples/cpp/logistic_regression.cpp

@@ -1,10 +1,15 @@
 // Copyright © 2023 Apple Inc.
 
 #include <chrono>
-#include <cmath>
+#include <functional>
 #include <iostream>
+#include <ratio>
 
-#include "mlx/mlx.h"
+#include "mlx/array.h"
+#include "mlx/ops.h"
+#include "mlx/random.h"
+#include "mlx/transforms.h"
+#include "mlx/utils.h"
 #include "timer.h"
 
 /**

examples/cpp/metal_capture.cpp

@@ -1,9 +1,14 @@
 // Copyright © 2024 Apple Inc.
 
-#include <cassert>
 #include <iostream>
-
-#include "mlx/mlx.h"
+#include <string>
+
+#include "mlx/backend/metal/metal.h"
+#include "mlx/device.h"
+#include "mlx/dtype.h"
+#include "mlx/ops.h"
+#include "mlx/stream.h"
+#include "mlx/utils.h"
 
 namespace mx = mlx::core;
 

examples/cpp/tutorial.cpp

@@ -1,9 +1,15 @@
 // Copyright © 2023 Apple Inc.
 
+#include <stddef.h>
 #include <cassert>
+#include <functional>
 #include <iostream>
 
-#include "mlx/mlx.h"
+#include "mlx/array.h"
+#include "mlx/dtype.h"
+#include "mlx/ops.h"
+#include "mlx/transforms.h"
+#include "mlx/utils.h"
 
 namespace mx = mlx::core;
 

mlx/array.cpp

@@ -3,10 +3,13 @@
 #include <unordered_map>
 
 #include "mlx/array.h"
+#include "mlx/device.h"
 #include "mlx/ops.h"
 #include "mlx/primitives.h"
+#include "mlx/stream.h"
 #include "mlx/transforms.h"
 #include "mlx/transforms_impl.h"
+#include "mlx/types/complex.h"
 
 namespace mlx::core {
 

mlx/array.h

@@ -1,16 +1,31 @@
 // Copyright © 2023 Apple Inc.
 #pragma once
 
+#include <arm_fp16.h>
+#include <stddef.h>
 #include <algorithm>
+#include <complex>
 #include <cstdint>
 #include <functional>
+#include <initializer_list>
+#include <iterator>
 #include <memory>
+#include <stdexcept>
+#include <type_traits>
+#include <utility>
 #include <vector>
 
 #include "mlx/allocator.h"
 #include "mlx/dtype.h"
 #include "mlx/event.h"
 #include "mlx/small_vector.h"
+#include "mlx/types/half_types.h"
+
+namespace mlx {
+namespace core {
+struct complex64_t;
+} // namespace core
+} // namespace mlx
 
 namespace mlx::core {
 

mlx/backend/common/broadcasting.cpp

@@ -1,6 +1,7 @@
 // Copyright © 2024 Apple Inc.
 
-#include "mlx/backend/common/utils.h"
+#include "mlx/allocator.h"
+#include "mlx/array.h"
 
 namespace mlx::core {
 

mlx/backend/common/common.cpp

@@ -1,9 +1,23 @@
 // Copyright © 2024 Apple Inc.
+#include <arm_fp16.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <algorithm>
 #include <cassert>
+#include <stdexcept>
+#include <utility>
+#include <vector>
 
+#include "mlx/allocator.h"
+#include "mlx/array.h"
 #include "mlx/backend/common/broadcasting.h"
 #include "mlx/backend/common/utils.h"
+#include "mlx/dtype.h"
 #include "mlx/primitives.h"
+#include "mlx/small_vector.h"
+#include "mlx/types/bf16.h"
+#include "mlx/types/complex.h"
+#include "mlx/types/half_types.h"
 
 namespace mlx::core {
 

mlx/backend/common/compiled.cpp

@@ -1,7 +1,18 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include "mlx/backend/common/compiled.h"
+
+#include <arm_fp16.h>
+#include <stdint.h>
+#include <algorithm>
+#include <sstream>
+#include <stdexcept>
+#include <utility>
+
 #include "mlx/backend/common/utils.h"
+#include "mlx/dtype.h"
+#include "mlx/types/complex.h"
+#include "mlx/types/half_types.h"
 #include "mlx/utils.h"
 
 namespace mlx::core {

mlx/backend/common/compiled.h

@@ -1,12 +1,28 @@
 // Copyright © 2023-2024 Apple Inc.
 #pragma once
 
+#include <stddef.h>
 #include <functional>
 #include <iomanip>
-
+#include <limits>
+#include <ostream>
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <typeinfo>
+#include <vector>
+
+#include "mlx/allocator.h"
 #include "mlx/array.h"
 #include "mlx/primitives.h"
 
+namespace mlx {
+namespace core {
+struct Dtype;
+struct _MLX_BFloat16;
+} // namespace core
+} // namespace mlx
+
 namespace mlx::core {
 
 inline bool is_static_cast(const Primitive& p) {

mlx/backend/common/load.cpp

@@ -1,10 +1,18 @@
 // Copyright © 2023 Apple Inc.
 
-#include <algorithm>
+#include <stddef.h>
+#include <stdint.h>
+#include <future>
+#include <memory>
+#include <type_traits>
 #include <utility>
 
+#include "mlx/allocator.h"
+#include "mlx/array.h"
+#include "mlx/io/load.h"
 #include "mlx/primitives.h"
 #include "mlx/scheduler.h"
+#include "mlx/threadpool.h"
 
 namespace {
 

mlx/backend/common/reduce.cpp

@@ -2,6 +2,9 @@
 
 #include "mlx/backend/common/reduce.h"
 
+#include <stdint.h>
+#include <algorithm>
+
 namespace mlx::core {
 
 std::pair<Shape, Strides> shapes_without_reduction_axes(

mlx/backend/common/reduce.h

@@ -2,6 +2,10 @@
 
 #pragma once
 
+#include <utility>
+#include <vector>
+
+#include "mlx/array.h"
 #include "mlx/backend/common/utils.h"
 
 namespace mlx::core {

mlx/backend/common/slicing.cpp

@@ -1,5 +1,14 @@
 // Copyright © 2024 Apple Inc.
 
+#include <stddef.h>
+#include <stdint.h>
+#include <memory>
+#include <sstream>
+#include <stdexcept>
+#include <tuple>
+
+#include "mlx/allocator.h"
+#include "mlx/array.h"
 #include "mlx/backend/common/utils.h"
 
 namespace mlx::core {

mlx/backend/common/slicing.h

@@ -2,6 +2,9 @@
 
 #pragma once
 
+#include <stdint.h>
+#include <tuple>
+
 #include "mlx/array.h"
 
 namespace mlx::core {

mlx/backend/common/utils.cpp

@@ -1,6 +1,7 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include <dlfcn.h>
+#include <stdexcept>
 
 #include "mlx/backend/common/utils.h"
 

mlx/backend/common/utils.h

@@ -2,11 +2,19 @@
 
 #pragma once
 
+#include <_stdlib.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <algorithm>
 #include <filesystem>
+#include <iterator>
+#include <limits>
 #include <tuple>
+#include <utility>
 #include <vector>
 
 #include "mlx/array.h"
+#include "mlx/small_vector.h"
 
 namespace mlx::core {
 

mlx/backend/cpu/arg_reduce.cpp

@@ -1,10 +1,21 @@
 // Copyright © 2023 Apple Inc.
 
+#include <arm_fp16.h>
+#include <stdint.h>
 #include <cassert>
+#include <functional>
+#include <vector>
 
+#include "mlx/allocator.h"
+#include "mlx/array.h"
 #include "mlx/backend/common/utils.h"
 #include "mlx/backend/cpu/encoder.h"
+#include "mlx/dtype.h"
 #include "mlx/primitives.h"
+#include "mlx/stream.h"
+#include "mlx/types/bf16.h"
+#include "mlx/types/complex.h"
+#include "mlx/types/half_types.h"
 
 namespace mlx::core {
 

mlx/backend/cpu/compiled.cpp

@@ -1,22 +1,44 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include <dlfcn.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <exception>
 #include <filesystem>
 #include <fstream>
+#include <functional>
+#include <iomanip>
 #include <list>
 #include <mutex>
 #include <shared_mutex>
-
-#include <fmt/format.h>
-
+#include <sstream>
+#include <stdexcept>
+#include <string>
+#include <string_view>
+#include <tuple>
+#include <type_traits>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "fmt/core.h"
+#include "mlx/array.h"
 #include "mlx/backend/common/compiled.h"
 #include "mlx/backend/cpu/compiled_preamble.h"
 #include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/jit_compiler.h"
-#include "mlx/device.h"
 #include "mlx/graph_utils.h"
+#include "mlx/primitives.h"
+#include "mlx/small_vector.h"
+#include "mlx/stream.h"
 #include "mlx/version.h"
 
+namespace mlx {
+namespace core {
+struct Device;
+} // namespace core
+} // namespace mlx
+
 namespace mlx::core {
 
 struct CompilerCache {