Update to support elementwiseAffine for RMS, reflect for pad.

Author: liuliu

WORKSPACE

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

nnc/ModelAddons.swift

@@ -270,15 +270,16 @@ extension ModelIOConvertible {
 public final class Pad: Model {
   public enum Mode {
     case zero
-    case replication
-    public static let clampToEdge: Mode = .replication
+    case replicate
+    case reflect
+    public static let clampToEdge: Mode = .replicate
   }
   required init(_ model: OpaquePointer) {
     super.init(model)
   }
 
   public init(
-    _ mode: Mode, begin: TensorShape, end: TensorShape, name: String = ""
+    _ mode: Mode, begin: TensorShape = [], end: TensorShape = [], name: String = ""
   ) {
     var begin = begin.dims
     var end = end.dims
@@ -290,8 +291,10 @@ public final class Pad: Model {
         switch mode {
         case .zero:
           type = Int32(CCV_NNC_PAD_ZERO)
-        case .replication:
+        case .replicate:
           type = Int32(CCV_NNC_PAD_REPLICATE)
+        case .reflect:
+          type = Int32(CCV_NNC_PAD_REFLECT)
         }
         return ccv_cnnp_pad(type, begin, end, name)!
       }
@@ -308,7 +311,7 @@ extension ModelIOConvertible {
    *   - begin: The beginning pad for each dimension.
    *   - end: The end pad for each dimension.
    */
-  public func padded(_ mode: Pad.Mode, begin: TensorShape, end: TensorShape) -> Model.IO {
+  public func padded(_ mode: Pad.Mode, begin: TensorShape = [], end: TensorShape = []) -> Model.IO {
     return Pad(mode, begin: begin, end: end)(self)
   }
 }
@@ -715,11 +718,15 @@ public final class RMSNorm: Model {
     super.init(model)
   }
 
-  public init(epsilon: Float, axis: [Int], trainable: Bool? = nil, name: String = "") {
+  public init(
+    epsilon: Float, axis: [Int], elementwiseAffine: Bool = true, 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),
+        epsilon, axis32, Int32(axis.count), elementwiseAffine ? 1 : 0,
+        trainable == true ? 1 : (trainable == false ? 0 : -1),
         name))
   }