More shape type (#1705)

* more shape type * fix
2024-12-19 08:08:20 -08:00
parent f17536af9c
commit e03f0372b1
38 changed files with 260 additions and 258 deletions
@@ -181,7 +181,7 @@ void init_fast(nb::module_& parent_module) {
        return nb::cpp_function(
            [kernel = std::move(kernel)](
                const std::vector<ScalarOrArray>& inputs_,
-                const std::vector<std::vector<int>>& output_shapes,
+                const std::vector<mx::Shape>& output_shapes,
                const std::vector<mx::Dtype>& output_dtypes,
                std::tuple<int, int, int> grid,
                std::tuple<int, int, int> threadgroup,