Add {Tensor,tensor,ndarray}.square() example.

Author: Iainmon

lib/Autograd.chpl

@@ -278,6 +278,17 @@ record reluOp : serializable {
     proc spec : GradOpSpec do return new dict(("operation","ReLU"));
 }
 
+record squareOp : serializable {
+    var input: shared BaseTensorResource(?);
+
+    proc children do return (input,);
+
+    proc forward() do
+        return input.array.square();
+
+    proc spec : GradOpSpec do return new dict(("operation","Square"));
+}
+
 record expOp : serializable {
     var input: shared BaseTensorResource(?);
 

lib/DynamicTensor.chpl

@@ -244,6 +244,15 @@ proc dynamicTensor.relu(): dynamicTensor(eltType) {
     return new dynamicTensor(eltType);
 }
 
+proc dynamicTensor.square(): dynamicTensor(eltType) {
+    for param rank in 1..maxRank {
+        if this.checkRank(rank) then
+            return this.forceRank(rank).square().eraseRank();
+    }
+    halt("Could not determine rank in dynamicTensor.square.");
+    return new dynamicTensor(eltType);
+}
+
 proc dynamicTensor.gelu(): dynamicTensor(eltType) {
     for param rank in 1..maxRank {
         if this.checkRank(rank) then

lib/NDArray.chpl

@@ -582,6 +582,18 @@ record ndarray : serializable {
         return rl;
     }
 
+    inline proc square() {
+        const ref thisData = data;
+        const dom = this.domain;
+        var rl = new ndarray(dom,eltType);
+        ref rlD = rl.data;
+        forall i in dom.every() {
+            const x = thisData[i];
+            rlD[i] = x * x;
+        }
+        return rl;
+    }
+
     inline proc gelu() {
         const ref thisData = data;
         const dom = this.domain;

lib/StaticTensor.chpl

@@ -177,6 +177,11 @@ proc staticTensor.relu() {
     return tensorFromCtx(rank,eltType,ctx);
 }
 
+proc staticTensor.square() {
+    var ctx = new squareOp(meta);
+    return tensorFromCtx(rank,eltType,ctx);
+}
+
 proc staticTensor.gelu() {
     var t = new staticTensor(rank,eltType);
     on this.device {