fix scatter + test (#1202)

* fix scatter + test

* fix test warnings

* fix metal validation

mlx/backend/metal/jit/indexing.h

@@ -38,12 +38,24 @@ constexpr std::string_view scatter_kernels = R"(
     device mlx_atomic<{1}>* out [[buffer(2)]],
     const constant int* out_shape [[buffer(3)]],
     const constant size_t* out_strides [[buffer(4)]],
-    const constant size_t& upd_size [[buffer(5)]],
+    const constant size_t& out_ndim [[buffer(5)]],
+    const constant int* upd_shape [[buffer(6)]],
+    const constant size_t& upd_ndim [[buffer(7)]],
+    const constant size_t& upd_size [[buffer(8)]],
     {5}
     uint2 gid [[thread_position_in_grid]]) {{
   const array<const device {2}*, {4}> idx_buffers = {{ {6} }};
   return scatter_1d_index_impl<{1}, {2}, {3}, {4}>(
-      updates, out, out_shape, out_strides, upd_size, idx_buffers, gid);
+      updates,
+      out,
+      out_shape,
+      out_strides,
+      out_ndim,
+      upd_shape,
+      upd_ndim,
+      upd_size,
+      idx_buffers,
+      gid);
 }}
 
 [[kernel]] void scatter{0}_{4}(