diff --git a/WORKSPACE b/WORKSPACE index 6bae752c416..48d92b6047e 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -5,9 +5,9 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") git_repository( name = "ccv", - commit = "8e4a5a05201c1cf9d966a283c014b182c14368b5", + commit = "a916bbf1b7d05f3d1573bd2f9a0108c07c098a7c", remote = "https://github.com/liuliu/ccv.git", - shallow_since = "1702616811 -0500", + shallow_since = "1702668857 -0500", ) load("@ccv//config:ccv.bzl", "ccv_deps", "ccv_setting") diff --git a/deps.bzl b/deps.bzl index a416e72c1bb..48252962046 100644 --- a/deps.bzl +++ b/deps.bzl @@ -17,8 +17,8 @@ def s4nnc_deps(): git_repository, name = "ccv", remote = "https://github.com/liuliu/ccv.git", - commit = "8e4a5a05201c1cf9d966a283c014b182c14368b5", - shallow_since = "1702616811 -0500", + commit = "a916bbf1b7d05f3d1573bd2f9a0108c07c098a7c", + shallow_since = "1702668857 -0500", ) _maybe( diff --git a/nnc/ModelAddons.swift b/nnc/ModelAddons.swift index 13005f8bba6..a92b113d5a5 100644 --- a/nnc/ModelAddons.swift +++ b/nnc/ModelAddons.swift @@ -56,6 +56,33 @@ public final class Mul: Model { } } +/// Div two inputs together. It will not do broadcast. +public final class Div: Model { + required init(_ model: OpaquePointer) { + super.init(model) + } + + public init(reciprocal: Bool = false, name: String = "") { + super.init(ccv_cnnp_div(reciprocal ? 1 : 0, name)) + } + + public func callAsFunction( + _ left: T, _ right: T, streamContext: StreamContext? = nil + ) -> T { + let outputs = self(inputs: left, right, streamContext: streamContext) + return T(outputs[0]) + } +} + +extension ModelIOConvertible { + /** + * Compute the reciprocal for a model IO. + */ + public func reciprocal() -> Model.IO { + return Div(reciprocal: true)(self) + } +} + /// Matrix-multiplication over two inputs. public final class Matmul: Model { required init(_ model: OpaquePointer) { diff --git a/nnc/Operators.swift b/nnc/Operators.swift index c681d757ac6..42d791ef6c9 100644 --- a/nnc/Operators.swift +++ b/nnc/Operators.swift @@ -1,6 +1,6 @@ -infix operator .*: MultiplicationPrecedence -infix operator ./: MultiplicationPrecedence -infix operator .+: AdditionPrecedence +infix operator .* : MultiplicationPrecedence +infix operator ./ : MultiplicationPrecedence +infix operator .+ : AdditionPrecedence // Element-wise addition public func .+ (left: T, right: T) -> T { @@ -16,6 +16,10 @@ public func ./ (left: T, right: T) -> T { return Functional.div(left: left, right: right) } +public func ./ (left: ModelIOConvertible, right: ModelIOConvertible) -> Model.IO { + return Div(reciprocal: false)(left, right) +} + // Broadcast element-wise multiplication public func .* (left: T, right: T) -> T { return Functional.mul(left: left, right: right) @@ -84,6 +88,22 @@ public func + (left: ModelIOConvertible, right: Float) -> Model.IO { return Add()(left, Scalar(value: right)(left)) } +public func / (left: Float, right: ModelIOConvertible) -> Model.IO { + if left == 1 { + return Div(reciprocal: true)(right) + } else { + return left * Div(reciprocal: true)(right) + } +} + +public func / (left: ModelIOConvertible, right: Float) -> Model.IO { + if right == 1 { + return left.io + } else { + return Scalmul(1.0 / right)(left) + } +} + // Broadcast element-wise subtraction. public func - (left: T, right: T) -> T { return Functional.add(left: left, right: right, rightScalar: -1) diff --git a/test/model.swift b/test/model.swift index 4e316448552..8a7232481be 100644 --- a/test/model.swift +++ b/test/model.swift @@ -138,6 +138,27 @@ final class ModelTests: XCTestCase { XCTAssertEqual(tv3.rawValue[0], 1.1 * 2.2 + 3.1, accuracy: 1e-5) } + func testModelDiv() throws { + let dynamicGraph = DynamicGraph() + + func DivRec() -> Model { + let i0 = Input() + let i1 = Input() + let i2 = i0 ./ i1 + let i3 = Input() + let i4 = 0.5 / i3 + return Model([i0, i1, i3], [i2, i4]) + } + + let div = DivRec() + let tv0 = dynamicGraph.variable(Tensor([1.1], .CPU, .C(1))) + let tv1 = dynamicGraph.variable(Tensor([2.2], .CPU, .C(1))) + let tv2 = dynamicGraph.variable(Tensor([0.2], .CPU, .C(1))) + let tv3s = div(inputs: tv0, tv1, tv2).map { $0.as(of: Float32.self) } + XCTAssertEqual(tv3s[0].rawValue[0], 1.1 / 2.2, accuracy: 1e-5) + XCTAssertEqual(tv3s[1].rawValue[0], 0.5 / 0.2, accuracy: 1e-5) + } + func testModelScaledDotProductAttention() throws { let dynamicGraph = DynamicGraph() let q = dynamicGraph.variable(Tensor(.CPU, .NHWC(1, 10, 8, 20))) @@ -347,6 +368,7 @@ final class ModelTests: XCTestCase { ("testSequential", testSequential), ("testModelWithScalar", testModelWithScalar), ("testModelWithParameter", testModelWithParameter), + ("testModelDiv", testModelDiv), ("testModelScaledDotProductAttention", testModelScaledDotProductAttention), ("testCustomModel", testCustomModel), ]