Update library/layers.yaml to follow TC coding style

Author: nicolasvasilache

tensor_comprehensions/library/layers.yaml

@@ -26,225 +26,223 @@
 - name: indexing
   lang: |
     def indexing(float(H, W) input, int32(L) index) -> (output) {{
-      output(l, w) = input(index(l), w) where l in 0:{L}
+        output(l, w) = input(index(l), w) where l in 0:{L}
     }}
 
 - name: lookup_table
   lang: |
     def lookup_table(float(B, R) LUT, int32(B, N) I) -> (O) {
-      O(b, n) +=! LUT(I(b, n), r)
+        O(b, n) +=! LUT(I(b, n), r_r)
     }
 
 - name: matmul
   lang: |
-    def matmul(float(M,N) A, float(N,K) B) -> (output) {
-      output(i, j) +=! A(i, kk) * B(kk, j)
+    def matmul(float(M, K) A, float(K, N) B) -> (C) {
+        C(m, n) +=! A(m, r_k) * B(r_k, n)
     }
   grad: |
-    def matmul_grad(float(M,N) A, float(N,K), float(M,K) O_grad) -> (A_grad, B_grad){
-      A_grad(i, j) +=! O_grad(i, kk) * B(j, kk)
-      B_grad(i, j) +=! O_grad(kk, j) * A(kk, i)
+    def matmul_grad(float(M,K) A, float(K,N) B, float(M,N) g_C) -> (g_A, g_B){
+        g_A(m, k) +=! g_C(  m, r_n) * B(  k, r_n)
+        g_B(k, n) +=! g_C(r_m,   n) * A(r_m,   k)
     }
 
 - name: batch_matmul
   lang: |
     def batch_matmul(float(B, N, M) X, float(B, M, K) Y) -> (Z) {
-      Z(b, n, k) +=! X(b, n, mm) * Y(b, mm, k)
+        Z(b, n, k) +=! X(b, n, r_m) * Y(b, r_m, k)
     }
 
 - name: transpose
   lang: |
     def transpose(float(N, C, H, W) I) -> (O) {
-      O(c, n, w, h) = I(n, c, h, w)
+        O(c, n, w, h) = I(n, c, h, w)
     }
 
 - name: avgpool
   lang: |
     def avgpool(float(B, C, H, W) input) -> (output) {{
-      output(b, c, h, w) +=! input(b, c, h * {sH} + kh, w * {sW} + kw) / ({kH} * {kW}) where kh in 0:{kH}, kw in 0:{kW}
+        output(b, c, h, w) +=! input(b, c, h * {sH} + r_kh, w * {sW} + r_kw) / ({kH} * {kW})
+            where r_kh in 0:{kH}, r_kw in 0:{kW}
     }}
 
 - name: maxpool
   lang: |
     def maxpool(float(B, C, H, W) input) -> (output) {{
-      output(b, c, h, w) max=! input(b, c, h * {sH} + kh, w * {sW} + kw) where kh in 0:{kH}, kw in 0:{kW}
+        output(b, c, h, w) max=! input(b, c, h * {sH} + r_kh, w * {sW} + r_kw)
+            where r_kh in 0:{kH}, r_kw in 0:{kW}
     }}
 
 - name: scale
   lang: |
     def scale(float(M, N) I) -> (O) {{
-      O(m, n) = I(m, n) * {s}
+        O(m, n) = I(m, n) * {s}
     }}
 
 - name: sigmoid
   lang: |
     def sigmoid(float(N, C, H, W) I) -> (O) {
-      O(n, c, h, w) = 1 / (1 + exp(-I(n, c, h, w)))
+        O(n, c, h, w) = 1 / (1 + exp(-I(n, c, h, w)))
     }
 
 - name: softmax
   lang: |
     def softmax(float(N, D) I) -> (O, expsum, maxVal) {
-      maxVal(n) max= I(n, d)
-      expsum(n) +=! exp(I(n, d) - maxVal(n))
-      O(n, d) = exp(I(n, d)) / expsum(n)
+        maxVal(n) max=      I(n, d)
+          expsum(n) +=! exp(I(n, d) - maxVal(n))
+             O(n, d) =  exp(I(n, d)) / expsum(n)
     }
 
 - name: Tanh
   lang: |
     def Tanh(float(M) I) -> (O) {
-      O(m) = tanh(I(m))
+        O(m) = tanh(I(m))
     }
 
 - name: tensordot
   lang: |
     def tensordot(float(N, C1, C2, H, W) I0, float(N, C2, C3, H, W) I1) -> (O) {
-      O(n, c1, c3, h, w) +=! I0(n, c1, c2, h, w) * I1(n, c2, c3, h, w)
+        O(n, c1, c3, h, w) +=! I0(n, c1, r_c2, h, w) * I1(n, r_c2, c3, h, w)
     }
 
 - name: fully_connected
   lang: |
     def fully_connected(float(B, M) I, float(N, M) W1, float(N) B1) -> (O1) {
-      O1(b, n) +=! I(b, m) * W1(n, m)
-      O1(b, n) = O1(b, n) + B1(n)
+        O1(b, n) +=! I(b, r_m) * W1(n, r_m)
+        O1(b, n) = O1(b, n) + B1(n)
     }
 
 - name: relu
   lang: |
     def relu(float(B, M) I) -> (O1){
-      O1(b, m) = fmax(I(b, m), 0)
+        O1(b, m) = fmax(I(b, m), 0)
     }
 
 - name: fcrelu
   lang: |
-    def fcrelu(float(B, M) I, float(N, M) W1, float(N) B1) -> (O1){
-      O1(b, n) +=! I(b, m) * W1(n, m)
-      O1(b, n) = O1(b, n) + B1(n)
-      O1(b, n) = fmax(O1(b, n), 0)
+    def fcrelu(float(B,M) I, float(N,M) W1, float(N) B1) -> (O1){
+        O1(b, n) +=! I(b, r_m) * W1(n, r_m)
+        O1(b, n)  = O1(b,   n) + B1(n)
+        O1(b, n)  = fmax(O1(b, n), 0)
     }
 
 - name: cast
   lang: |
-    def cast(float(M, N) A) -> (int32(M, N) O1) {{
-      O1(m, n) = int32(A(m, n) + {constant})
+    def cast(float(M,N) A) -> (int32(M,N) O1) {{
+        O1(m, n) = int32(A(m, n) + {constant})
     }}
 
 - name: concat
   lang: |
     def concat(float(M, N) A, float(M, N) B) -> (O1) {
-      O1(n, i, m) = i == 0 ? A(m, n) : B(m, n) where i in 0:2
+        O1(n, i, m) = i == 0 ? A(m, n) : B(m, n) where i in 0:2
     }
 
 - name: convolution
   lang: |
     def convolution(float(N, C, H, W) I, float(M, C, KH, KW) W1, float(M) B) -> (O) {
-      O(n, m, h, w) +=! I(n, c, h + kh, w + kw) * W1(m, c, kh, kw)
-      O(n, m, h, w) = O(n, m, h, w) + B(m)
+        O(n, m, h, w) +=! I(n, r_c, h + r_kh, w + r_kw) * W1(m, r_c, r_kh, r_kw)
+        O(n, m, h, w)  =  O(n, m, h, w) + B(m)
     }
 
 - name: convolution_strided
   lang: |
     def convolution_strided(float(N, C, H, W) I, float(M, C, KH, KW) W1, float(M) B) -> (O) {{
-      O(n, m, h, w) +=! I(n, c, {sh} * h + kh, {sw} * w + kw) * W1(m, c, kh, kw)
-      O(n, m, h, w) = O(n, m, h, w) + B(m)
+        O(n, m, h, w) +=! I(n, r_c, {sh} * h + r_kh, {sw} * w + r_kw) * W1(m, r_c, r_kh, r_kw)
+        O(n, m, h, w)  = O(n, m, h, w) + B(m)
     }}
   grad: |
-    def convolution_strided_grad(float(N, C, H, W) I, float(M, C, KH, KW) W1, float(N, M, H, W) O_grad)
-    -> (I_grad, W1_grad) {{
-      I_grad(n, c, h, w) +=! O_grad(n, m, {sh} * h - kh, {sw} * w - kw) * W1(m, c, kh, kw)
-      W1_grad(m, c, kh, kw) +=! O_grad(n, m, {sh} * h - kh, {sw} * w - kw) * I(n, c, h, w)
+    def convolution_strided_grad(float(N, C, H, W) I, float(M, C, KH, KW) W1, float(N, M, H, W) g_O) -> (g_I, g_W1) {{
+         g_I(n, c, h, w)   +=! g_O(n, r_m, {sh} *   h - r_kh, {sw} *   w - r_kw) * W1(r_m, c, r_kh, r_kw)
+        g_W1(m, c, kh, kw) +=! g_O(n,   m, {sh} * r_h -   kh, {sw} * r_w -   kw) *  I(r_n, c,  r_h,  r_w)
     }}
 
 - name: group_convolution
   lang: |
     def group_convolution(float(N, G, C, H, W) I, float(G, F, C, KH, KW) W1, float(G, F) B) -> (O) {
-      O(n, g, f, h, w) +=! I(n, g, c, h + kh, w + kw) * W1(g, f, c, kh, kw)
-      O(n, g, f, h, w) = O(n, g, f, h, w) + B(g, f)
+        O(n, g, f, h, w) +=! I(n, g, r_c, h + r_kh, w + r_kw) * W1(g, f, r_c, r_kh, r_kw)
+        O(n, g, f, h, w)  =  O(n, g, f, h, w) + B(g, f)
     }
 
 - name: group_convolution_strided
   lang: |
-    def group_convolution_strided(float(N, G, C, H, W) I, float(G, F, C, KH, KW) W1, float(G, F) B) -> (O)
-    {{
-      O(n, g, f, h, w) +=! I(n, g, c, {sh} * h + kh, {sw} * w + kw) * W1(g, f, c, kh, kw)
-      O(n, g, f, h, w) = O(n, g, f, h, w) + B(g, f)
+    def group_convolution_strided(float(N, G, C, H, W) I, float(G, F, C, KH, KW) W1, float(G, F) B) -> (O) {{
+        O(n, g, f, h, w) +=! I(n, g, r_c, {sh} * h + r_kh, {sw} * w + r_kw) * W1(g, f, r_c, r_kh, r_kw)
+        O(n, g, f, h, w) = O(n, g, f, h, w) + B(g, f)
     }}
 
 - name: copy2D
   lang: |
     def copy2D(float(M, N) I) -> (O) {
-      O(i, j) = I(i, j)
+        O(m, n) = I(m, n)
     }
 
 - name: copy
   lang: |
     def copy(float({dimParams}) I) -> (O) {{
-      O({dimIndices}) = I({dimIndices})
+        O({dimIndices}) = I({dimIndices})
     }}
 
 - name: cosine
   lang: |
     def cosine(float(M) I) -> (O) {
-      O(i) = cos(I(i))
+        O(i) = cos(I(i))
     }
 
 - name: cosine_similarity
   lang: |
     def cosine_similarity(float(M, N) I1, float(M, N) I2) -> (O, sumI1, sumI2) {{
-      sumI1(m) +=! I1(m, n) * I1(m, n)
-      sumI2(m) +=! I2(m, n) * I2(m, n)
-      O(m) +=! (I1(m, n) * I2(m, n)) / fmax(rsqrt(sumI1(m)) * sqrt(sumI2(m)), {eps})
+       sumI1(m) +=!  I1(m, r_n) * I1(m, r_n)
+       sumI2(m) +=!  I2(m, r_n) * I2(m, r_n)
+           O(m) +=! (I1(m, r_n) * I2(m, r_n)) / fmax(rsqrt(sumI1(m)) * sqrt(sumI2(m)), {eps})
     }}
 
 - name: add
   lang: |
     def add(float(N) A, float(N) B) -> (output) {
-      output(i) = A(i) + B(i)
+        output(n) = A(n) + B(n)
     }
 
 - name: abs
   lang: |
     def abs(float(M, N) A) -> (O1) {
-      O1(m, n) = fabs(A(m, n))
+        O1(m, n) = fabs(A(m, n))
     }
 
 - name: layernorm
   lang: |
     def layernorm(float(T, B, C) I) -> (O, mean, centered, var) {{
-      mean(t, b) +=! I(t, b, c) / C
-      centered(t, b, c) = I(t, b, c) - mean(t, b)
-      var(t, b) +=! centered(t, b, c) * centered(t, b, c)
-      var(t, b) = (var(t, b) + {eps}) / C
-      O(t, b, c) = centered(t, b, c) / rsqrt(var(t, b))
+              mean(t, b) +=! I(t, b, c) / C
+        centered(t, b, c) =  I(t, b, c) - mean(t, b)
+        var(t, b) +=! centered(t, b, c) * centered(t, b, c)
+        var(t, b)  =  (var(t, b) + {eps}) / C
+        O(t, b, c) =  centered(t, b, c) / rsqrt(var(t, b))
     }}
 
 - name: batchnorm
   lang: |
-    def batchnorm(float(N, C, H, W) I, float(C) rMeanIn, float(C) rVarIn)
+    def batchnorm(float(N,C,H,W) I, float(C) rMeanIn, float(C) rVarIn)
     -> (O, rMeanOut, rVarOut, mean, centered, variance, expectedVariance, normalizedOut)
     {{
-       mean(c) +=! I(nn, c, hh, ww)
-       mean(c)  = mean(c) / (N * H * W)
-       rMeanOut(c) = (1 - {momentum}) * rMeanIn(c) + {momentum} * mean(c)
-       centered(n, c, h, w) = I(n, c, h, w) - rMeanOut(c)
-       variance(n, c, h, w) = centered(n, c, h, w) * centered(n, c, h, w)
-       expectedVariance(c) +=! (variance(n, c, h, w) + {eps}) / (N * H * W)
-       rVarOut(c) = rsqrt(
-         (1 - {momentum}) * rVarIn(c) + {momentum} * expectedVariance(c))
-       O(n, c, h, w) = centered(n, c, h, w) * rVarOut(c)
-       normalizedOut(n, c, h, w) = O(n, c, h, w)
+        mean(c) +=! I(nn, c, hh, ww)
+        mean(c)  = mean(c) / (N * H * W)
+        rMeanOut(c) = (1 - {momentum}) * rMeanIn(c) + {momentum} * mean(c)
+        centered(n, c, h, w) =        I(n, c, h, w) - rMeanOut(c)
+        variance(n, c, h, w) = centered(n, c, h, w) * centered(n, c, h, w)
+        expectedVariance(c) +=! (variance(n, c, h, w) + {eps}) / (N * H * W)
+        rVarOut(c) = rsqrt((1 - {momentum}) * rVarIn(c) + {momentum} * expectedVariance(c))
+        O(n, c, h, w) = centered(n, c, h, w) * rVarOut(c)
+        normalizedOut(n, c, h, w) = O(n, c, h, w)
     }}
 
 - name: small_mobilenet
   lang: |
     def small_mobilenet(float(C1, H, W) I, float(C1, KH1, KW1) W1, float(C1) B1, float(C2, C1) W2, float(C2) B2)
-    -> (O1, O2)
-    {
-      O1(c1, h, w) +=! I(c1, h + kh, w + kw) * W1(c1, kh, kw)
-      O1(c1, h, w)  = O1(c1, h, w) + B1(c1)
-      O1(c1, h, w)  = fmax(O1(c1, h, w), 0)
-
-      O2(c2, h, w) +=! O1(c1, h, w) * W2(c2, c1)
-      O2(c2, h, w)  = O2(c2, h, w) + B2(c2)
-      O2(c2, h, w)  = fmax(O2(c2, h, w), 0)
+    -> (O1, O2) {
+        O1(c1, h, w) +=! I(c1, h + r_kh, w + r_kw) * W1(c1, r_kh, r_kw)
+        O1(c1, h, w)  = O1(c1,        h,        w) + B1(c1)
+        O1(c1, h, w)  = fmax(O1(c1, h, w), 0)
+
+        O2(c2, h, w) +=! O1(r_c1, h, w) * W2(c2, r_c1)
+        O2(c2, h, w)  =  O2(  c2, h, w) + B2(c2)
+        O2(c2, h, w)  = fmax(O2(c2, h, w), 0)
     }