Dynamic broadcasting for shapeless compile/export (#1722)

* working towards dynamic broadcast

* shapeless broadcast

* fix build + nits

* use broadcast arrays in quantize matmul

* some cleanup / consistency

* mend

* some comments

* add vjp, jvp for broadcast axes

mlx/backend/accelerate/primitives.cpp

@@ -32,6 +32,7 @@ DEFAULT(ArgSort)
 DEFAULT(AsStrided)
 DEFAULT(BlockMaskedMM)
 DEFAULT(Broadcast)
+DEFAULT(BroadcastAxes)
 DEFAULT(Ceil)
 DEFAULT(Concatenate)
 DEFAULT(Conjugate)

mlx/backend/common/common.cpp

@@ -42,9 +42,7 @@ void AsStrided::eval(const std::vector<array>& inputs, array& out) {
   return move_or_copy(in, out, strides_, flags, data_size, offset_);
 }
 
-void Broadcast::eval(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 1);
-  const auto& in = inputs[0];
+void broadcast(const array& in, array& out) {
   if (out.size() == 0) {
     out.set_data(nullptr);
     return;
@@ -61,6 +59,14 @@ void Broadcast::eval(const std::vector<array>& inputs, array& out) {
   move_or_copy(in, out, strides, flags, in.data_size());
 }
 
+void Broadcast::eval(const std::vector<array>& inputs, array& out) {
+  broadcast(inputs[0], out);
+}
+
+void BroadcastAxes::eval(const std::vector<array>& inputs, array& out) {
+  broadcast(inputs[0], out);
+}
+
 void Copy::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   move_or_copy(inputs[0], out);

mlx/backend/common/default_primitives.cpp

@@ -37,6 +37,7 @@ DEFAULT(ArgSort)
 DEFAULT(AsType)
 DEFAULT(AsStrided)
 DEFAULT(Broadcast)
+DEFAULT(BroadcastAxes)
 DEFAULT(BlockMaskedMM)
 DEFAULT(GatherMM)
 DEFAULT(GatherQMM)

mlx/backend/metal/compiled.cpp

@@ -1,6 +1,5 @@
 // Copyright © 2023-2024 Apple Inc.
 #include <fmt/format.h>
-#include <iostream> //TODO
 #include <sstream>
 
 #include "mlx/backend/common/compiled.h"

mlx/backend/metal/primitives.cpp

@@ -240,6 +240,10 @@ void Broadcast::eval_gpu(const std::vector<array>& inputs, array& out) {
   eval(inputs, out);
 }
 
+void BroadcastAxes::eval_gpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+
 void Concatenate::eval_gpu(const std::vector<array>& inputs, array& out) {
   concatenate_gpu(inputs, out, axis_, stream());
 }

mlx/backend/no_cpu/primitives.cpp

@@ -35,6 +35,7 @@ NO_CPU(AsStrided)
 NO_CPU(BitwiseBinary)
 NO_CPU(BlockMaskedMM)
 NO_CPU(Broadcast)
+NO_CPU(BroadcastAxes)
 NO_CPU(Ceil)
 NO_CPU(Cholesky)
 NO_CPU(Concatenate)

mlx/backend/no_metal/primitives.cpp

@@ -36,6 +36,7 @@ NO_GPU(AsStrided)
 NO_GPU(BitwiseBinary)
 NO_GPU(BlockMaskedMM)
 NO_GPU(Broadcast)
+NO_GPU(BroadcastAxes)
 NO_GPU(Ceil)
 NO_GPU_MULTI(Compiled)
 NO_GPU(Concatenate)