Expose RMSNorm and cmul.

WORKSPACE

@@ -5,9 +5,9 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
 git_repository(
     name = "ccv",
-    commit = "bb296cdcca8a6dca24934b675c3dd7bd12278245",
+    commit = "ff47b596600c501d1e6d7070a5e6454a90dce1be",
     remote = "https://github.com/liuliu/ccv.git",
-    shallow_since = "1703102984 -0500",
+    shallow_since = "1703182006 -0500",
 )
 
 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 = "bb296cdcca8a6dca24934b675c3dd7bd12278245",
-        shallow_since = "1703102984 -0500",
+        commit = "ff47b596600c501d1e6d7070a5e6454a90dce1be",
+        shallow_since = "1703182006 -0500",
     )
 
     _maybe(

nnc/FunctionalAddons.swift

@@ -308,6 +308,19 @@ extension Functional {
     return T(outputs[0])
   }
 
+  /// Complex number multiplication
+  public static func cmul<T: DynamicGraph.TensorGroup>(
+    left: T, right: T, streamContext: StreamContext? = nil
+  ) -> T {
+    var params = CmdParamsFactory.factory.newParams()
+    params.size.dim = (1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0)
+    let cmd = ccv_nnc_cmd(CCV_NNC_CMUL_FORWARD, nil, params, 0)
+    let outputs = exec(
+      cmd: cmd, hint: ccv_nnc_no_hint, inputs: left, right, outputSize: 1,
+      streamContext: streamContext)
+    return T(outputs[0])
+  }
+
   /// Make a copy.
   public static func copy<T: DynamicGraph.TensorGroup>(
     from: T, to: T, streamContext: StreamContext? = nil

nnc/ModelAddons.swift

@@ -56,6 +56,14 @@ public final class Mul: Model {
   }
 }
 
+extension Functional {
+  public static func mul(left: ModelIOConvertible, right: ModelIOConvertible, scalar: Float)
+    -> Model.IO
+  {
+    return Mul(scalar: scalar)(left, right)
+  }
+}
+
 /// Div two inputs together. It will not do broadcast.
 public final class Div: Model {
   required init(_ model: OpaquePointer) {
@@ -130,6 +138,39 @@ public final class Matmul: Model {
   }
 }
 
+extension Functional {
+  public static func matmul(
+    left: ModelIOConvertible, right: ModelIOConvertible, leftTranspose: (Int, Int),
+    rightTranspose: (Int, Int)
+  ) -> Model.IO {
+    return Matmul(transposeA: leftTranspose, transposeB: rightTranspose)(left, right)
+  }
+}
+
+/// Comlex number multiplication over two inputs.
+public final class Cmul: Model {
+  required init(_ model: OpaquePointer) {
+    super.init(model)
+  }
+
+  public init(name: String = "") {
+    super.init(ccv_cnnp_cmul(name))
+  }
+
+  public func callAsFunction<T: DynamicGraph.TensorGroup>(
+    _ left: T, _ right: T, streamContext: StreamContext? = nil
+  ) -> T {
+    let outputs = self(inputs: left, right, streamContext: streamContext)
+    return T(outputs[0])
+  }
+}
+
+extension Functional {
+  public static func cmul(left: ModelIOConvertible, right: ModelIOConvertible) -> Model.IO {
+    return Cmul()(left, right)
+  }
+}
+
 /// A linear layer model.
 public final class Dense: Model {
   required init(_ model: OpaquePointer) {
@@ -605,6 +646,28 @@ public final class GroupNorm: Model {
   }
 }
 
+/// RMSNorm model.
+public final class RMSNorm: Model {
+  required init(_ model: OpaquePointer) {
+    super.init(model)
+  }
+
+  public init(epsilon: Float, axis: [Int], trainable: Bool? = nil, name: String = "") {
+    let axis32: [Int32] = axis.map { Int32($0) }
+    super.init(
+      ccv_cnnp_rmsnorm(
+        epsilon, axis32, Int32(axis.count), trainable == true ? 1 : (trainable == false ? 0 : -1),
+        name))
+  }
+
+  public func callAsFunction<T: DynamicGraph.TensorGroup>(
+    _ input: T, streamContext: StreamContext? = nil
+  ) -> T {
+    let outputs = self(inputs: input, streamContext: streamContext)
+    return T(outputs[0])
+  }
+}
+
 /// Make the input tensor to be 1-D tensor (respecting N).
 public final class Flatten: Model {
   required init(_ model: OpaquePointer) {