Fix tf.pow for zero exponent (#1823)

BUG

This fixes tf.pow(0, 0), which in WebGL1 returns NaN but in
WebGL2 and JavaScript return 1.

Note that TensorFlow also defines pow(0,0) = 1.

Author: AndreasMadsen

src/backends/webgl/binaryop_gpu.ts

@@ -33,7 +33,7 @@ export const MUL = 'return a * b;';
 export const DIV = `
 if (b == 0.0) {
   return NAN;
-} 
+}
 if (a == b) {
   return 1.0;
 };
@@ -59,6 +59,9 @@ export const POW = `
 if(a < 0.0 && floor(b) < b){
   return NAN;
 }
+if (b == 0.0) {
+  return 1.0;
+}
 return (round(mod(b, 2.0)) != 1) ?
     pow(abs(a), b) : sign(a) * pow(abs(a), b);
 `;

src/backends/webgl/binaryop_packed_gpu.ts

@@ -55,7 +55,7 @@ export const DIV = `
   } else if(a.w == b.w) {
     result.w = 1.;
   }
-  
+
   return result;
 `;
 
@@ -88,6 +88,13 @@ export const POW = `
   vec4 multiplier = sign(a) * isModRound1 + (vec4(1.0) - isModRound1);
   vec4 result = multiplier * pow(abs(a), b);
 
+  // Ensure that a^0 = 1, including 0^0 = 1 as this correspond to TF and JS
+  bvec4 isExpZero = equal(b, vec4(0.0));
+  result.r = isExpZero.r ? 1.0 : result.r;
+  result.g = isExpZero.g ? 1.0 : result.g;
+  result.b = isExpZero.b ? 1.0 : result.b;
+  result.a = isExpZero.a ? 1.0 : result.a;
+
   vec4 isNaN = vec4(lessThan(a, vec4(0.0))) * vec4(lessThan(floor(b), b));
   ` +
     CHECK_NAN_SNIPPET + `

src/ops/arithmetic_test.ts

@@ -738,6 +738,14 @@ describeWithFlags('pow', ALL_ENVS, () => {
     expectArraysClose(await result.data(), [NaN, 27, NaN, 0], 0.05);
   });
 
+  it('exponent of 0 returns 1', async () => {
+    const a = tf.tensor1d([-2, -1, 0, 1, 2]);
+    const b = tf.scalar(0);
+
+    const result = tf.pow(a, b);
+    expectArraysClose(await result.data(), [1, 1, 1, 1, 1]);
+  });
+
   it('handles non int32 exponent param', async () => {
     const a = tf.tensor1d([2, 4]);
     const b = tf.tensor1d([.5, 1.2]);