Fix flaky TestVmap.test_vmap_masked_scatter (#3421)

mlx/backend/cuda/indexing.cpp

@@ -458,7 +458,7 @@ void MaskedScatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   }
 
   array mask_flat = flatten_in_eval(mask, 1, -1, s);
-  if (mask_flat.data<void>() != mask.data<void>()) {
+  if (gpu_ptr<void>(mask_flat) != gpu_ptr<void>(mask)) {
     encoder.add_temporary(mask_flat);
   }
   if (!mask_flat.flags().row_contiguous) {

mlx/backend/cuda/scan.cu

@@ -364,7 +364,7 @@ constexpr bool supports_scan_op() {
 }
 
 void scan_gpu_inplace(
-    array in,
+    const array& in,
     array& out,
     Scan::ReduceType reduce_type,
     int axis,

mlx/backend/gpu/scan.h

@@ -6,7 +6,7 @@
 namespace mlx::core {
 
 void scan_gpu_inplace(
-    array in,
+    const array& in,
     array& out,
     Scan::ReduceType reduce_type,
     int axis,

mlx/backend/metal/scan.cpp

@@ -13,7 +13,7 @@
 namespace mlx::core {
 
 void scan_gpu_inplace(
-    array in,
+    const array& in,
     array& out,
     Scan::ReduceType reduce_type,
     int axis,