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

python/src/ops.cpp

@@ -2799,6 +2799,27 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The output array with the new shape.
       )pbdoc");
+  m.def(
+      "broadcast_arrays",
+      [](const nb::args& args, mx::StreamOrDevice s) {
+        return broadcast_arrays(nb::cast<std::vector<mx::array>>(args), s);
+      },
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def broadcast_arrays(*arrays: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, ...]"),
+      R"pbdoc(
+        Broadcast arrays against one another.
+
+        The broadcasting semantics are the same as Numpy.
+
+        Args:
+            *arrays (array): The input arrays.
+
+        Returns:
+            tuple(array): The output arrays with the broadcasted shape.
+      )pbdoc");
   m.def(
       "softmax",
       [](const mx::array& a,
@@ -3853,8 +3874,8 @@ void init_ops(nb::module_& m) {
 
         Args:
             file (file, str): Path to file to which the arrays are saved.
-            args (arrays): Arrays to be saved.
-            kwargs (arrays): Arrays to be saved. Each array will be saved
+            *args (arrays): Arrays to be saved.
+            **kwargs (arrays): Arrays to be saved. Each array will be saved
               with the associated keyword as the output file name.
       )pbdoc");
   m.def(