diff --git a/mlx/CMakeLists.txt b/mlx/CMakeLists.txt
index f6d1a4aa2d..52206d34e8 100644
--- a/mlx/CMakeLists.txt
+++ b/mlx/CMakeLists.txt
@@ -1,6 +1,7 @@
 target_sources(
   mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/compile.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/dtype.cpp
diff --git a/mlx/allocator.cpp b/mlx/allocator.cpp
new file mode 100644
index 0000000000..b91d690990
--- /dev/null
+++ b/mlx/allocator.cpp
@@ -0,0 +1,53 @@
+// Copyright © 2026 Apple Inc.
+
+#include <algorithm>
+#include <mutex>
+
+#include "mlx/allocator.h"
+
+namespace mlx::core::allocator {
+
+#ifndef NDEBUG
+
+bool Allocator::has_invalid_memory(const std::vector<void*>& ptrs) {
+  std::shared_lock lock(allocations_mutex_);
+  return std::any_of(ptrs.begin(), ptrs.end(), [this](void* ptr) {
+    if (!ptr) {
+      return false;
+    }
+    auto it = allocations_.upper_bound(ptr);
+    if (it == allocations_.begin()) {
+      return true;
+    }
+    --it;
+    auto offset = static_cast<char*>(ptr) - static_cast<char*>(it->first);
+    return static_cast<size_t>(offset) > it->second;
+  });
+}
+
+void Allocator::track_allocation(void* data_ptr, size_t size) {
+  if (data_ptr && size > 0) {
+    std::unique_lock lock(allocations_mutex_);
+    allocations_[data_ptr] = size;
+  }
+}
+
+void Allocator::untrack_allocation(void* data_ptr) {
+  if (data_ptr) {
+    std::unique_lock lock(allocations_mutex_);
+    allocations_.erase(data_ptr);
+  }
+}
+
+void Allocator::swap_allocation(void* data_ptr_a, void* data_ptr_b) {
+  std::unique_lock lock(allocations_mutex_);
+  auto it = allocations_.find(data_ptr_a);
+  if (it != allocations_.end()) {
+    allocations_[data_ptr_b] = it->second;
+    allocations_.erase(it);
+  }
+}
+
+#endif // NDEBUG
+
+} // namespace mlx::core::allocator
diff --git a/mlx/allocator.h b/mlx/allocator.h
index 6ef8599455..82bc87401d 100644
--- a/mlx/allocator.h
+++ b/mlx/allocator.h
@@ -6,6 +6,12 @@
 #include <functional>
 #include <utility>
 
+#ifndef NDEBUG
+#include <map>
+#include <shared_mutex>
+#include <vector>
+#endif
+
 #include "mlx/api.h"
 
 namespace mlx::core::allocator {
@@ -52,12 +58,29 @@ class MLX_API Allocator {
   };
   virtual void release(Buffer buffer) {}
 
+#ifndef NDEBUG
+  bool has_invalid_memory(const std::vector<void*>& ptrs);
+#endif
+
   Allocator() = default;
   Allocator(const Allocator& other) = delete;
   Allocator(Allocator&& other) = delete;
   Allocator& operator=(const Allocator& other) = delete;
   Allocator& operator=(Allocator&& other) = delete;
   virtual ~Allocator() = default;
+
+ protected:
+#ifndef NDEBUG
+  void track_allocation(void* data_ptr, size_t size);
+  void untrack_allocation(void* data_ptr);
+  void swap_allocation(void* data_ptr_a, void* data_ptr_b);
+#endif
+
+ private:
+#ifndef NDEBUG
+  std::map<void*, size_t> allocations_;
+  std::shared_mutex allocations_mutex_;
+#endif
 };
 
 MLX_API Allocator& allocator();
diff --git a/mlx/backend/common/common.cpp b/mlx/backend/common/common.cpp
index 1e0692973c..81dc5c1a07 100644
--- a/mlx/backend/common/common.cpp
+++ b/mlx/backend/common/common.cpp
@@ -41,7 +41,7 @@ void AsStrided::eval(const array& in, array& out) {
   flags.row_contiguous = row_contiguous;
   flags.col_contiguous = col_contiguous;
 
-  return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
+  out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
 }
 
 void Broadcast::eval(const std::vector<array>& inputs, array& out) {
diff --git a/mlx/backend/cpu/primitives.cpp b/mlx/backend/cpu/primitives.cpp
index 1633e0a647..406e08eb34 100644
--- a/mlx/backend/cpu/primitives.cpp
+++ b/mlx/backend/cpu/primitives.cpp
@@ -91,7 +91,7 @@ void AsStrided::eval_cpu(const std::vector<array>& inputs, array& out) {
     eval(inputs[0], out);
   } else {
     auto copy = contiguous_copy_cpu(inputs[0], stream());
-    cpu::get_command_encoder(stream()).add_temporary(copy);
+    // Donated to out.
     eval(copy, out);
   }
 }
diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp
index bb262d0e63..4c44ed3826 100644
--- a/mlx/backend/cuda/allocator.cpp
+++ b/mlx/backend/cuda/allocator.cpp
@@ -251,6 +251,9 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
   if (buf->device >= 0 && buf->device != device) {
     move_to_unified_memory(*buf, stream);
   }
+#ifndef NDEBUG
+  track_allocation(buf->data, size);
+#endif
   return Buffer{buf};
 }
 
@@ -268,6 +271,10 @@ void CudaAllocator::free(Buffer buffer) {
     return;
   }
 
+#ifndef NDEBUG
+  untrack_allocation(buf->data);
+#endif
+
   std::unique_lock lock(mutex_);
   active_memory_ -= buf->size;
   if (get_cache_memory() < max_pool_size_) {
@@ -301,6 +308,9 @@ void CudaAllocator::move_to_unified_memory(
     CHECK_CUDA_ERROR(cudaMemcpy(data, buf.data, buf.size, kind));
     free_async(buf);
   }
+#ifndef NDEBUG
+  swap_allocation(buf.data, data);
+#endif
   buf.data = data;
   buf.device = -1;
 }
diff --git a/mlx/backend/cuda/conv/gemm_grouped_conv.cu b/mlx/backend/cuda/conv/gemm_grouped_conv.cu
index 82d32aae72..23055c2032 100644
--- a/mlx/backend/cuda/conv/gemm_grouped_conv.cu
+++ b/mlx/backend/cuda/conv/gemm_grouped_conv.cu
@@ -163,6 +163,7 @@ void gemm_grouped_conv_nd(
   wt_view.copy_shared_buffer(
       wt, {wt.strides(0), 1, C_per_group}, wt.flags(), wt.size());
   array wt_reshaped = contiguous_copy_gpu(wt_view, s);
+  encoder.add_temporary(wt_reshaped);
 
   // Batch with size of groups.
   Shape batch_shape{params.groups};
diff --git a/mlx/backend/cuda/copy.cu b/mlx/backend/cuda/copy.cu
index d3d67e6578..a19d15550d 100644
--- a/mlx/backend/cuda/copy.cu
+++ b/mlx/backend/cuda/copy.cu
@@ -109,7 +109,11 @@ void fill_gpu(const array& in, array& out, const Stream& s) {
   copy_contiguous(encoder, CopyType::Scalar, in, out, 0, 0);
 }
 
-void reshape_gpu(const array& in, array& out, Stream s) {
+void reshape_gpu(
+    const array& in,
+    array& out,
+    bool add_out_to_temporary_on_allocation,
+    Stream s) {
   auto [copy_necessary, out_strides] = prepare_reshape(in, out);
   if (copy_necessary) {
     auto& encoder = cu::get_command_encoder(s);
@@ -124,6 +128,9 @@ void reshape_gpu(const array& in, array& out, Stream s) {
         0,
         CopyType::General,
         s);
+    if (add_out_to_temporary_on_allocation) {
+      encoder.add_temporary(out);
+    }
   } else {
     shared_buffer_reshape(in, out_strides, out);
   }
diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp
index 52363446b7..c432ab1cb4 100644
--- a/mlx/backend/cuda/device.cpp
+++ b/mlx/backend/cuda/device.cpp
@@ -245,16 +245,25 @@ void CommandEncoder::set_input_array(const array& arr) {
   bytes_in_graph_ += arr.data_size();
   auto id = reinterpret_cast<std::uintptr_t>(arr.buffer().ptr());
   active_deps_.push_back(id);
+#ifndef NDEBUG
+  if (arr.data_shared_ptr()->is_owned()) {
+    registered_data_ptrs_.push_back(gpu_ptr<void>(arr));
+  }
+#endif
 }
 
 void CommandEncoder::set_output_array(const array& arr) {
   if (!use_cuda_graphs()) {
     return;
   }
-
   auto id = reinterpret_cast<std::uintptr_t>(arr.buffer().ptr());
   active_deps_.push_back(id);
   active_outputs_.push_back(id);
+#ifndef NDEBUG
+  if (arr.data_shared_ptr()->is_owned()) {
+    registered_data_ptrs_.push_back(gpu_ptr<void>(arr));
+  }
+#endif
 }
 
 void CommandEncoder::add_kernel_node_raw(
@@ -514,6 +523,19 @@ void CommandEncoder::synchronize() {
   f.wait();
 }
 
+#ifndef NDEBUG
+void CommandEncoder::verify_graph(const char* primitive_name) {
+  if (allocator::allocator().has_invalid_memory(registered_data_ptrs_)) {
+    throw std::runtime_error(
+        fmt::format(
+            "There are invalid memory registered when calling {}::eval_gpu, "
+            "which is a bug that corrupts the graph, and usually caused by "
+            "temporary arrays not registered with \"add_temporary\" method.",
+            primitive_name));
+  }
+}
+#endif
+
 void CommandEncoder::commit_impl() {
   if (!temporaries_.empty()) {
     add_completed_handler([temporaries = std::move(temporaries_)]() {});
@@ -607,6 +629,9 @@ void CommandEncoder::clear_graph_state() {
   active_outputs_.clear();
   concurrent_nodes_.clear();
   roots_.clear();
+#ifndef NDEBUG
+  registered_data_ptrs_.clear();
+#endif
   is_graph_updatable_ = true;
 }
 
diff --git a/mlx/backend/cuda/device.h b/mlx/backend/cuda/device.h
index 6bc64d8df7..56e7cee2f5 100644
--- a/mlx/backend/cuda/device.h
+++ b/mlx/backend/cuda/device.h
@@ -112,6 +112,10 @@ class MLX_API CommandEncoder {
   void add_completed_handler(std::function<void()> task);
   bool needs_commit();
   void commit();
+  void synchronize();
+#ifndef NDEBUG
+  void verify_graph(const char* primitive_name);
+#endif
 
   Device& device() {
     return device_;
@@ -121,9 +125,6 @@ class MLX_API CommandEncoder {
     return stream_;
   }
 
-  // Wait until kernels and completion handlers are finished
-  void synchronize();
-
  private:
   cudaGraphNode_t add_kernel_node_raw(const cudaKernelNodeParams& params);
   CUgraphNode add_kernel_node_raw(const CUDA_KERNEL_NODE_PARAMS& params);
@@ -147,7 +148,7 @@ class MLX_API CommandEncoder {
   CudaStream stream_;
   CudaGraph graph_;
   std::shared_ptr<Worker> worker_;
-  int node_count_{0};
+
   bool in_concurrent_{false};
   std::vector<cudaGraphNode_t> from_nodes_;
   std::vector<cudaGraphNode_t> to_nodes_;
@@ -160,11 +161,17 @@ class MLX_API CommandEncoder {
   std::vector<std::uintptr_t> active_outputs_;
   std::unordered_map<std::uintptr_t, GraphNode> node_map_;
   std::vector<cudaGraphNode_t> roots_;
-  size_t bytes_in_graph_{0};
+#ifndef NDEBUG
+  std::vector<void*> registered_data_ptrs_;
+#endif
   bool is_graph_updatable_{true};
-  bool force_single_root_{false};
+
+  int node_count_{0};
+  size_t bytes_in_graph_{0};
   int max_ops_per_graph_;
   int max_mb_per_graph_;
+
+  bool force_single_root_;
 };
 
 class MLX_API Device {
diff --git a/mlx/backend/cuda/eval.cpp b/mlx/backend/cuda/eval.cpp
index 2862f38223..9f8d24e033 100644
--- a/mlx/backend/cuda/eval.cpp
+++ b/mlx/backend/cuda/eval.cpp
@@ -70,6 +70,10 @@ void eval(array& arr) {
     encoder.add_temporary(s);
   }
 
+#ifndef NDEBUG
+  encoder.verify_graph(arr.primitive().name());
+#endif
+
   if (encoder.needs_commit()) {
     scheduler::notify_new_task(stream);
     encoder.add_completed_handler(
diff --git a/mlx/backend/cuda/indexing.cpp b/mlx/backend/cuda/indexing.cpp
index fbdba986c6..bc62c3a278 100644
--- a/mlx/backend/cuda/indexing.cpp
+++ b/mlx/backend/cuda/indexing.cpp
@@ -479,9 +479,6 @@ void MaskedScatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   array mask_flat = (mask.ndim() == 1)
       ? reshape_in_eval(mask, {mask.shape(0), 1}, s)
       : flatten_in_eval(mask, 1, -1, s);
-  if (gpu_ptr<void>(mask_flat) != gpu_ptr<void>(mask)) {
-    encoder.add_temporary(mask_flat);
-  }
   if (!mask_flat.flags().row_contiguous) {
     mask_flat = contiguous_copy_gpu(mask_flat, s);
     encoder.add_temporary(mask_flat);
diff --git a/mlx/backend/gpu/copy.cpp b/mlx/backend/gpu/copy.cpp
index 699dadd1a6..2084667347 100644
--- a/mlx/backend/gpu/copy.cpp
+++ b/mlx/backend/gpu/copy.cpp
@@ -56,7 +56,7 @@ array flatten_in_eval(const array& x, int start_axis, int end_axis, Stream s) {
 
 array reshape_in_eval(const array& x, Shape shape, Stream s) {
   array out(std::move(shape), x.dtype(), nullptr, {});
-  reshape_gpu(x, out, s);
+  reshape_gpu(x, out, true, s);
   return out;
 }
 
diff --git a/mlx/backend/gpu/copy.h b/mlx/backend/gpu/copy.h
index 8ebcd84648..807ff0e1ef 100644
--- a/mlx/backend/gpu/copy.h
+++ b/mlx/backend/gpu/copy.h
@@ -48,7 +48,11 @@ void fill_gpu(const array& val, array& out, const Stream& s);
 array contiguous_copy_gpu(const array& arr, const Stream& s);
 
 // Copy data from |in| and transpose to |out|'s shape.
-void reshape_gpu(const array& in, array& out, Stream s);
+void reshape_gpu(
+    const array& in,
+    array& out,
+    bool add_out_to_temporary_on_allocation,
+    Stream s);
 
 // Like the normal ops but safe to call in eval_gpu.
 array flatten_in_eval(const array& x, int start_axis, int end_axis, Stream s);
diff --git a/mlx/backend/gpu/primitives.cpp b/mlx/backend/gpu/primitives.cpp
index 6df6ab3e09..5470372ad8 100644
--- a/mlx/backend/gpu/primitives.cpp
+++ b/mlx/backend/gpu/primitives.cpp
@@ -30,11 +30,7 @@ void AsStrided::eval_gpu(const std::vector<array>& inputs, array& out) {
     eval(inputs[0], out);
   } else {
     auto copy = contiguous_copy_gpu(inputs[0], stream());
-#if defined(MLX_USE_CUDA)
-    cu::get_command_encoder(stream()).add_temporary(copy);
-#else
-    metal::get_command_encoder(stream()).add_temporary(copy);
-#endif
+    // Donated to out.
     eval(copy, out);
   }
 }
@@ -183,7 +179,7 @@ void Full::eval_gpu(const std::vector<array>& inputs, array& out) {
 
 void Flatten::eval_gpu(const std::vector<array>& inputs, array& out) {
   MLX_PROFILER_RANGE("Flatten::eval_gpu");
-  reshape_gpu(inputs[0], out, stream());
+  reshape_gpu(inputs[0], out, false, stream());
 }
 
 void NumberOfElements::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -209,7 +205,7 @@ void Pad::eval_gpu(const std::vector<array>& inputs, array& out) {
 
 void Reshape::eval_gpu(const std::vector<array>& inputs, array& out) {
   MLX_PROFILER_RANGE("Reshape::eval_gpu");
-  reshape_gpu(inputs[0], out, stream());
+  reshape_gpu(inputs[0], out, false, stream());
 }
 
 void Split::eval_gpu(
@@ -248,7 +244,7 @@ void Transpose::eval_gpu(const std::vector<array>& inputs, array& out) {
 
 void Unflatten::eval_gpu(const std::vector<array>& inputs, array& out) {
   MLX_PROFILER_RANGE("Unflatten::eval_gpu");
-  reshape_gpu(inputs[0], out, stream());
+  reshape_gpu(inputs[0], out, false, stream());
 }
 
 void View::eval_gpu(const std::vector<array>& inputs, array& out) {
diff --git a/mlx/backend/metal/allocator.cpp b/mlx/backend/metal/allocator.cpp
index 903cce2187..ff12293b79 100644
--- a/mlx/backend/metal/allocator.cpp
+++ b/mlx/backend/metal/allocator.cpp
@@ -172,7 +172,11 @@ Buffer MetalAllocator::malloc(size_t size) {
         get_cache_memory() - max_pool_size_);
   }
 
-  return Buffer{static_cast<void*>(buf)};
+  Buffer buffer(static_cast<void*>(buf));
+#ifndef NDEBUG
+  track_allocation(buffer.raw_ptr(), size);
+#endif
+  return buffer;
 }
 
 void MetalAllocator::clear_cache() {
@@ -185,6 +189,11 @@ void MetalAllocator::free(Buffer buffer) {
   if (buf == nullptr) {
     return;
   }
+
+#ifndef NDEBUG
+  untrack_allocation(buffer.raw_ptr());
+#endif
+
   std::unique_lock lk(mutex_);
   active_memory_ -= buf->length();
   if (get_cache_memory() < max_pool_size_) {
diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp
index 59acdc6ae9..d883fdb997 100644
--- a/mlx/backend/metal/conv.cpp
+++ b/mlx/backend/metal/conv.cpp
@@ -974,7 +974,7 @@ void stride_two_conv_transpose_3D_gpu(
     array phase_out({M, O}, out.dtype(), nullptr, {});
     phase_out.set_data(allocator::malloc(phase_out.nbytes()));
 
-    std::vector<array> gemm_copies = {in_matrix, wt_phase};
+    std::vector<array> gemm_copies;
     steel_matmul(
         s,
         d,
@@ -1018,7 +1018,6 @@ void stride_two_conv_transpose_3D_gpu(
     copy_gpu_inplace(phase_out_nd, out_phase_view, CopyType::GeneralGeneral, s);
     copies.push_back(phase_out);
     copies.push_back(phase_out_nd);
-    copies.push_back(out_phase_view);
   }
 }
 
diff --git a/mlx/backend/metal/copy.cpp b/mlx/backend/metal/copy.cpp
index 97c5fccbca..cce820dfdb 100644
--- a/mlx/backend/metal/copy.cpp
+++ b/mlx/backend/metal/copy.cpp
@@ -213,7 +213,11 @@ void fill_gpu(const array& val, array& out, const Stream& s) {
   compute_encoder.dispatch_threads(grid_dims, group_dims);
 }
 
-void reshape_gpu(const array& in, array& out, Stream s) {
+void reshape_gpu(
+    const array& in,
+    array& out,
+    bool add_out_to_temporary_on_allocation,
+    Stream s) {
   auto [copy_necessary, out_strides] = prepare_reshape(in, out);
   if (copy_necessary) {
     out.set_data(allocator::malloc(out.nbytes()));
@@ -227,6 +231,9 @@ void reshape_gpu(const array& in, array& out, Stream s) {
         0,
         CopyType::General,
         s);
+    if (add_out_to_temporary_on_allocation) {
+      metal::get_command_encoder(s).add_temporary(out);
+    }
   } else {
     shared_buffer_reshape(in, out_strides, out);
   }
diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp
index 493e8e3153..302a556638 100644
--- a/mlx/backend/metal/device.cpp
+++ b/mlx/backend/metal/device.cpp
@@ -365,6 +365,12 @@ void CommandEncoder::set_input_array(
   if (all_inputs_.insert(a.buffer().ptr()).second) {
     buffer_sizes_ += a.data_size();
   }
+#ifndef NDEBUG
+  if (a.data_shared_ptr()->is_owned()) {
+    registered_data_ptrs_.push_back(a.buffer().raw_ptr());
+  }
+#endif
+
   auto r_buf = static_cast<MTL::Resource*>(const_cast<void*>(a.buffer().ptr()));
   next_inputs_.insert(r_buf);
   needs_barrier_ =
@@ -384,6 +390,11 @@ void CommandEncoder::set_output_array(
 
 void CommandEncoder::register_output_array(const array& a) {
   all_outputs_.insert(a.buffer().ptr());
+#ifndef NDEBUG
+  if (a.data_shared_ptr()->is_owned()) {
+    registered_data_ptrs_.push_back(a.buffer().raw_ptr());
+  }
+#endif
 
   auto buf = static_cast<MTL::Resource*>(const_cast<void*>(a.buffer().ptr()));
   if (concurrent_) {
@@ -520,6 +531,9 @@ void CommandEncoder::end_encoding(bool wait_in_flight) {
   next_outputs_.clear();
   concurrent_outputs_.clear();
   all_inputs_.clear();
+#ifndef NDEBUG
+  registered_data_ptrs_.clear();
+#endif
 }
 
 void CommandEncoder::signal_event(Event event, uint64_t value) {
@@ -599,6 +613,44 @@ void CommandEncoder::synchronize() {
   }
 }
 
+#ifndef NDEBUG
+void CommandEncoder::verify_graph(
+    const char* primitive_name,
+    const std::vector<array>& inputs,
+    const std::vector<array>& outputs) {
+  if (exiting_) {
+    return;
+  }
+  if (allocator::allocator().has_invalid_memory(registered_data_ptrs_)) {
+    throw std::runtime_error(
+        fmt::format(
+            "There are invalid memory registered when calling {}::eval_gpu, "
+            "which is a bug that corrupts the graph, and usually caused by "
+            "temporary arrays not registered with \"add_temporary\" method.",
+            primitive_name));
+  }
+  for (const auto& temp : temporaries_) {
+    auto pred = [&temp](auto& arr) {
+      return arr.data_shared_ptr() == temp.data_shared_ptr();
+    };
+    if (std::find_if(inputs.begin(), inputs.end(), pred) != inputs.end()) {
+      throw std::runtime_error(
+          fmt::format(
+              "An input of array is registered as temporary when calling "
+              "{}::eval_gpu, which is a bug that corrupts the graph.",
+              primitive_name));
+    }
+    if (std::find_if(outputs.begin(), outputs.end(), pred) != outputs.end()) {
+      throw std::runtime_error(
+          fmt::format(
+              "A output of array is registered as temporary when calling "
+              "{}::eval_gpu, which is a bug that corrupts the graph.",
+              primitive_name));
+    }
+  }
+}
+#endif
+
 MTL::ComputeCommandEncoder* CommandEncoder::get_command_encoder() {
   if (!encoder_) {
     error_.check();
diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h
index 305a35a772..2f7b93652b 100644
--- a/mlx/backend/metal/device.h
+++ b/mlx/backend/metal/device.h
@@ -100,6 +100,12 @@ class MLX_API CommandEncoder {
   bool needs_commit() const;
   void commit(std::function<void()> completion = nullptr);
   void synchronize();
+#ifndef NDEBUG
+  void verify_graph(
+      const char* primitive_name,
+      const std::vector<array>& inputs,
+      const std::vector<array>& outputs);
+#endif
 
   MTL::CommandBuffer* get_command_buffer() const {
     return buffer_.get();
@@ -143,6 +149,9 @@ class MLX_API CommandEncoder {
   std::unordered_set<MTL::Resource*> concurrent_outputs_;
   std::unordered_set<const void*> all_inputs_;
   std::unordered_set<const void*> all_outputs_;
+#ifndef NDEBUG
+  std::vector<void*> registered_data_ptrs_;
+#endif
 
   // A map of prior command encoder outputs to their corresponding fence.
   std::unordered_map<const void*, NS::SharedPtr<MTL::Fence>> prev_ce_outputs_;
diff --git a/mlx/backend/metal/eval.cpp b/mlx/backend/metal/eval.cpp
index 43a6fb9932..b53feaea42 100644
--- a/mlx/backend/metal/eval.cpp
+++ b/mlx/backend/metal/eval.cpp
@@ -56,6 +56,10 @@ void eval(array& arr) {
     buffers.erase(it);
   }
 
+#ifndef NDEBUG
+  encoder.verify_graph(arr.primitive().name(), arr.inputs(), outputs);
+#endif
+
   if (encoder.needs_commit()) {
     encoder.end_encoding();
     scheduler::notify_new_task(s);
diff --git a/mlx/backend/metal/indexing.cpp b/mlx/backend/metal/indexing.cpp
index 2a74576664..8f4c46a179 100644
--- a/mlx/backend/metal/indexing.cpp
+++ b/mlx/backend/metal/indexing.cpp
@@ -658,10 +658,6 @@ void MaskedScatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   array mask_flat = (mask.ndim() == 1)
       ? reshape_in_eval(mask, {mask.shape(0), 1}, s)
       : flatten_in_eval(mask, 1, -1, s);
-  if (mask_flat.data<void>() != mask.data<void>()) {
-    compute_encoder.add_temporary(mask_flat);
-  }
-
   if (!mask_flat.flags().row_contiguous) {
     mask_flat = contiguous_copy_gpu(mask_flat, s);
     compute_encoder.add_temporary(mask_flat);
diff --git a/python/tests/test_device.py b/python/tests/test_device.py
index c0ed9ca900..bef788bb1d 100644
--- a/python/tests/test_device.py
+++ b/python/tests/test_device.py
@@ -134,9 +134,13 @@ def test_stream_context_exited_on_another_thread(self):
 
         def on_worker():
             def gen():
-                with mx.stream(mx.new_stream(mx.default_device())):
-                    yield 1
-                    yield 2
+                try:
+                    with mx.stream(mx.new_stream(mx.default_device())):
+                        yield 1
+                        yield 2
+                except Exception:
+                    # Exception is ignored on purpose.
+                    pass
 
             g = gen()
             next(g)