license import { KernelConfig , KernelFunc , Pow } from '@tensorflow/tfjs-core' ; import { CHECK_NAN_SNIPPET } from '../binaryop_packed_gpu' ; import { binaryKernelFunc } from '../kernel_utils/kernel_funcs_utils' ; 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);
` ; const POW_PACKED = `
  // isModRound1 has 1 for components with round(mod(b, 2.0)) == 1, 0 otherwise.
  vec4 isModRound1 = vec4(equal(round(mod(b, 2.0)), ivec4(1)));
  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 + `
  return result;
` ; export const pow = binaryKernelFunc ( { opSnippet : POW , packedOpSnippet : POW_PACKED } ) ; export const powConfig = { kernelName : Pow , backendName : 'webgl' , kernelFunc : pow as { } as KernelFunc } ;