Add Debug model.

WORKSPACE

@@ -5,9 +5,9 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
 git_repository(
     name = "ccv",
-    commit = "d2622a324f45c243c69fe8c2019eeaca9162ca0a",
+    commit = "ef7bd537aea035ab04e4f7d0d817f56f7b0492fc",
     remote = "https://github.com/liuliu/ccv.git",
-    shallow_since = "1726462252 -0400",
+    shallow_since = "1727501268 -0400",
 )
 
 load("@ccv//config:ccv.bzl", "ccv_deps", "ccv_setting")

deps.bzl

@@ -17,8 +17,8 @@ def s4nnc_deps():
         git_repository,
         name = "ccv",
         remote = "https://github.com/liuliu/ccv.git",
-        commit = "d2622a324f45c243c69fe8c2019eeaca9162ca0a",
-        shallow_since = "1726462252 -0400",
+        commit = "ef7bd537aea035ab04e4f7d0d817f56f7b0492fc",
+        shallow_since = "1727501268 -0400",
     )
 
     _maybe(

nnc/ModelAddons.swift

@@ -1498,6 +1498,40 @@ public final class ScaledDotProductAttention: Model {
   }
 }
 
+/// Debug model.
+public final class Debug: Model {
+  required init(_ model: OpaquePointer) {
+    super.init(model)
+  }
+
+  public init(
+    name: String = "",
+    _ callback: @escaping ([AnyTensor?], StreamContext?) -> Void
+  ) {
+    let underlying = Wrapped(callback)
+    super.init(
+      ccv_cnnp_debug(
+        { inputs, inputSize, stream, context in
+          let callback = Unmanaged<Wrapped<([AnyTensor?], StreamContext?) -> Void>>.fromOpaque(
+            context!
+          ).takeUnretainedValue().value
+          callback(
+            (0..<Int(inputSize)).map {
+              inputs![$0].map { AnyTensorStorage($0, selfOwned: false).toAnyTensor() }
+            }, stream.map { StreamContext(stream: $0, selfOwned: false) })
+        }, Unmanaged.passRetained(underlying).toOpaque(),
+        { context in
+          guard let context = context else { return }
+          Unmanaged<Wrapped<([AnyTensor?], StreamContext?) -> Void>>.fromOpaque(context).release()
+        },
+        { context in
+          guard let context = context else { return nil }
+          return Unmanaged<Wrapped<([AnyTensor?], StreamContext?) -> Void>>.fromOpaque(context)
+            .retain().toOpaque()
+        }, name))
+  }
+}
+
 /// Custom model.
 public final class CustomModel: Model {
   required init(_ model: OpaquePointer) {

nnc/StreamContext.swift

@@ -23,8 +23,14 @@ public final class StreamContext {
     }
   }
 
+  let selfOwned: Bool
   let _stream: OpaquePointer
 
+  init(stream: OpaquePointer, selfOwned: Bool) {
+    _stream = stream
+    self.selfOwned = selfOwned
+  }
+
   /**
    * Create a new stream context.
    *
@@ -39,6 +45,7 @@ public final class StreamContext {
       type = Int32((ordinal << 8) | CCV_STREAM_CONTEXT_GPU)
     }
     _stream = ccv_nnc_stream_context_new(type)!
+    selfOwned = true
   }
 
   /**
@@ -69,6 +76,7 @@ public final class StreamContext {
   }
 
   deinit {
+    guard selfOwned else { return }
     ccv_nnc_stream_context_free(_stream)
   }
 }

test/model.swift

@@ -171,6 +171,33 @@ final class ModelTests: XCTestCase {
     XCTAssertEqual(tv3s[1].rawValue[0], 0.5 / 0.2, accuracy: 1e-5)
   }
 
+  func testModelDebug() throws {
+    let dynamicGraph = DynamicGraph()
+
+    var value: Float = 0
+    func DivRec() -> Model {
+      let i0 = Input()
+      let i1 = Input()
+      let i2 = i0 ./ i1
+      let debugi2 =
+        (Debug { tensors, _ in
+          let tensor = Tensor<Float>(tensors[0]!)
+          value = tensor[0]
+        })(i2)
+      let i3 = Input()
+      let i4 = debugi2 ./ i3
+      return Model([i0, i1, i3], [i4])
+    }
+
+    let div = DivRec()
+    let tv0 = dynamicGraph.variable(Tensor<Float32>([1.1], .CPU, .C(1)))
+    let tv1 = dynamicGraph.variable(Tensor<Float32>([2.2], .CPU, .C(1)))
+    let tv2 = dynamicGraph.variable(Tensor<Float32>([0.2], .CPU, .C(1)))
+    let tv3s = div(inputs: tv0, tv1, tv2).map { $0.as(of: Float32.self) }
+    XCTAssertEqual(tv3s[0].rawValue[0], 0.5 / 0.2, accuracy: 1e-5)
+    XCTAssertEqual(value, 0.5, accuracy: 1e-5)
+  }
+
   func testModelScaledDotProductAttention() throws {
     let dynamicGraph = DynamicGraph()
     let q = dynamicGraph.variable(Tensor<Float32>(.CPU, .NHWC(1, 10, 8, 20)))
@@ -458,6 +485,7 @@ final class ModelTests: XCTestCase {
     ("testModelWithScalar", testModelWithScalar),
     ("testModelWithParameter", testModelWithParameter),
     ("testModelDiv", testModelDiv),
+    ("testModelDebug", testModelDebug),
     ("testModelScaledDotProductAttention", testModelScaledDotProductAttention),
     ("testCustomModel", testCustomModel),
     ("testModelShareWeights", testModelShareWeights),