Refactored PR Comment Issues, Plus Pulling in updates to tfjs master repository

WORKSPACE

@@ -50,7 +50,7 @@ load("@rules_nodejs//nodejs:repositories.bzl", "nodejs_register_toolchains")
 
 nodejs_register_toolchains(
     name = "nodejs",
-    node_version = "16.13.2",
+    node_version = "18.7.0",
 )
 
 # Install the yarn tool

package.json

@@ -58,7 +58,7 @@
     "rollup": "^2.46.0",
     "rollup-plugin-sourcemaps": "^0.6.3",
     "rollup-plugin-terser": "^7.0.2",
-    "rollup-plugin-visualizer": "~3.3.2",
+    "rollup-plugin-visualizer": "~5.8.2",
     "seedrandom": "^3.0.5",
     "shelljs": "~0.8.5",
     "string_decoder": "^1.3.0",

tfjs-backend-webgl/src/binaryop_packed_gpu.ts

@@ -21,11 +21,11 @@ import {GPGPUProgram, useShapeUniforms} from './gpgpu_math';
 import {getChannels} from './packing_util';
 import {getCoordsDataType} from './shader_compiler';
 
-export const CHECK_NAN_SNIPPET = `
-  result.r = isNaN.r > 0. ? NAN : result.r;
-  result.g = isNaN.g > 0. ? NAN : result.g;
-  result.b = isNaN.b > 0. ? NAN : result.b;
-  result.a = isNaN.a > 0. ? NAN : result.a;
+export const CHECK_NAN_SNIPPET_PACKED = `
+  result.r = isNaN.r ? NAN : result.r;
+  result.g = isNaN.g ? NAN : result.g;
+  result.b = isNaN.b ? NAN : result.b;
+  result.a = isNaN.a ? NAN : result.a;
 `;
 
 export const ELU_DER = `

tfjs-backend-webgl/src/flags_webgl.ts

@@ -268,3 +268,9 @@ ENV.registerFlag('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE', () => Infinity);
  * problem: https://github.com/tensorflow/tfjs/issues/6775.
  */
 ENV.registerFlag('WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE', () => false);
+
+/**
+ * Whether to use the customized isnan. It's only useful for webgl2 since webgl1
+ * doesn't have the builtin isnan.
+ */
+ENV.registerFlag('WEBGL2_ISNAN_CUSTOM', () => false);

tfjs-backend-webgl/src/glsl_version.ts

@@ -62,7 +62,7 @@ export function getGlslDifferences(): GLSL {
     //   - fraction = anything except all 0 bits (since all 0 bits represents
     //   infinity).
     // https://en.wikipedia.org/wiki/IEEE_754-1985#Representation_of_non-numbers
-    defineSpecialNaN = `
+    defineSpecialNaN = env().getBool('WEBGL2_ISNAN_CUSTOM') ? `
       bool isnan_custom(float val) {
         uint floatToUint = floatBitsToUint(val);
         return (floatToUint & 0x7fffffffu) > 0x7f800000u;
@@ -74,7 +74,8 @@ export function getGlslDifferences(): GLSL {
       }
 
       #define isnan(value) isnan_custom(value)
-    `;
+    ` :
+                                                              '';
     // In webgl 2 we do not need to specify a custom isinf so there is no
     // need for a special INFINITY constant.
     defineSpecialInf = ``;

tfjs-backend-webgl/src/kernel_utils/kernel_funcs_utils.ts

@@ -32,18 +32,6 @@ import {SimpleBinaryKernelImplCPU, SimpleUnaryKernelImplCPU} from './shared';
 
 export const CHECK_NAN_SNIPPET_UNARY = `if (isnan(x)) return x;`;
 
-export const CHECK_NAN_SNIPPET_BINARY = `
-  if (isnan(a)) return a;
-  if (isnan(b)) return b;
-`;
-
-export const CHECK_NAN_SNIPPET_BINARY_PACKED = `
-  result.r = isNaN.r > 0. ? NAN : result.r;
-  result.g = isNaN.g > 0. ? NAN : result.g;
-  result.b = isNaN.b > 0. ? NAN : result.b;
-  result.a = isNaN.a > 0. ? NAN : result.a;
-`;
-
 type UnaryKernelFuncConfig = {
   opSnippet: string,
   packedOpSnippet?: string,

tfjs-backend-webgl/src/kernels/Atan2.ts

@@ -17,18 +17,21 @@
 
 import {Atan2} from '@tensorflow/tfjs-core';
 import {KernelConfig} from '@tensorflow/tfjs-core';
+import {CHECK_NAN_SNIPPET} from '../binaryop_gpu';
+import {CHECK_NAN_SNIPPET_PACKED} from '../binaryop_packed_gpu';
+import {binaryKernelFunc} from '../kernel_utils/kernel_funcs_utils';
 
-import {binaryKernelFunc, CHECK_NAN_SNIPPET_BINARY, CHECK_NAN_SNIPPET_BINARY_PACKED} from '../kernel_utils/kernel_funcs_utils';
-
-const ATAN2 = CHECK_NAN_SNIPPET_BINARY + `
+const ATAN2 = CHECK_NAN_SNIPPET + `
   return atan(a, b);
 `;
 
 const ATAN2_PACKED = `
   vec4 result = atan(a, b);
-  vec4 isNaN = min(vec4(isnan(a)) + vec4(isnan(b)), vec4(1.0));
+  bvec4 isNaNA = isnan(a);
+  bvec4 isNaNB = isnan(b);
+  bvec4 isNaN = bvec4(isNaNA.x || isNaNB.x, isNaNA.y || isNaNB.y, isNaNA.z || isNaNB.z, isNaNA.w || isNaNB.w);
   ` +
-    CHECK_NAN_SNIPPET_BINARY_PACKED + `
+    CHECK_NAN_SNIPPET_PACKED + `
   return result;
 `;
 

tfjs-backend-webgl/src/kernels/Maximum.ts

@@ -18,7 +18,7 @@
 import {KernelConfig, KernelFunc, Maximum} from '@tensorflow/tfjs-core';
 
 import {CHECK_NAN_SNIPPET} from '../binaryop_gpu';
-import {CHECK_NAN_SNIPPET as CHECK_NAN_SNIPPET_PACKED} from '../binaryop_packed_gpu';
+import {CHECK_NAN_SNIPPET_PACKED} from '../binaryop_packed_gpu';
 import {binaryKernelFunc} from '../kernel_utils/kernel_funcs_utils';
 import {maximumImplCPU} from '../kernel_utils/shared';
 
@@ -28,7 +28,9 @@ const MAXIMUM = CHECK_NAN_SNIPPET + `
 
 const MAXIMUM_PACKED = `
   vec4 result = vec4(max(a, b));
-  vec4 isNaN = min(vec4(isnan(a)) + vec4(isnan(b)), vec4(1.0));
+  bvec4 isNaNA = isnan(a);
+  bvec4 isNaNB = isnan(b);
+  bvec4 isNaN = bvec4(isNaNA.x || isNaNB.x, isNaNA.y || isNaNB.y, isNaNA.z || isNaNB.z, isNaNA.w || isNaNB.w);
   ` +
     CHECK_NAN_SNIPPET_PACKED + `
   return result;

tfjs-backend-webgl/src/kernels/Minimum.ts

@@ -18,7 +18,7 @@
 import {KernelConfig, KernelFunc, Minimum} from '@tensorflow/tfjs-core';
 
 import {CHECK_NAN_SNIPPET} from '../binaryop_gpu';
-import {CHECK_NAN_SNIPPET as CHECK_NAN_SNIPPET_PACKED} from '../binaryop_packed_gpu';
+import {CHECK_NAN_SNIPPET_PACKED} from '../binaryop_packed_gpu';
 import {binaryKernelFunc} from '../kernel_utils/kernel_funcs_utils';
 import {minimumImplCPU} from '../kernel_utils/shared';
 
@@ -28,7 +28,9 @@ const MINIMUM = CHECK_NAN_SNIPPET + `
 
 const MINIMUM_PACKED = `
   vec4 result = vec4(min(a, b));
-  vec4 isNaN = min(vec4(isnan(a)) + vec4(isnan(b)), vec4(1.0));
+  bvec4 isNaNA = isnan(a);
+  bvec4 isNaNB = isnan(b);
+  bvec4 isNaN = bvec4(isNaNA.x || isNaNB.x, isNaNA.y || isNaNB.y, isNaNA.z || isNaNB.z, isNaNA.w || isNaNB.w);
   ` +
     CHECK_NAN_SNIPPET_PACKED + `
   return result;

tfjs-backend-webgl/src/kernels/Mod.ts

@@ -17,17 +17,17 @@
 
 import {KernelConfig, KernelFunc, Mod} from '@tensorflow/tfjs-core';
 
-import {CHECK_NAN_SNIPPET} from '../binaryop_packed_gpu';
+import {CHECK_NAN_SNIPPET_PACKED} from '../binaryop_packed_gpu';
 import {binaryKernelFunc} from '../kernel_utils/kernel_funcs_utils';
 
 const MOD = `if (b == 0.0) return NAN;
   return mod(a, b);`;
 
 const MOD_PACKED = `
   vec4 result = mod(a, b);
-  vec4 isNaN = vec4(equal(b, vec4(0.0)));
+  bvec4 isNaN = equal(b, vec4(0.0));
   ` +
-    CHECK_NAN_SNIPPET + `
+    CHECK_NAN_SNIPPET_PACKED + `
   return result;
 `;
 

tfjs-backend-webgl/src/kernels/Pow.ts

@@ -17,7 +17,7 @@
 
 import {KernelConfig, KernelFunc, Pow} from '@tensorflow/tfjs-core';
 
-import {CHECK_NAN_SNIPPET} from '../binaryop_packed_gpu';
+import {CHECK_NAN_SNIPPET_PACKED} from '../binaryop_packed_gpu';
 import {binaryKernelFunc} from '../kernel_utils/kernel_funcs_utils';
 
 const POW = `
@@ -44,9 +44,11 @@ const POW_PACKED = `
   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));
+  bvec4 isNaN1 = lessThan(a, vec4(0.0));
+  bvec4 isNaN2 = lessThan(floor(b), b);
+  bvec4 isNaN = bvec4(isNaN1.x && isNaN2.x, isNaN1.y && isNaN2.y, isNaN1.z && isNaN2.z, isNaN1.w && isNaN2.w);
   ` +
-    CHECK_NAN_SNIPPET + `
+    CHECK_NAN_SNIPPET_PACKED + `
   return result;
 `;
 

tfjs-converter/python/tensorflowjs/converters/fuse_prelu.py

@@ -87,7 +87,6 @@ def fuse_ops_for_prelu(input_graph_def):
       continue
 
     alpha_tensor_name = neg_alpha_op.name
-    _create_alpha_node(neg_alpha_op, updated_alpha)
 
     relu_neg_input_op = None
     for name in mul_op.input:
@@ -120,6 +119,7 @@ def fuse_ops_for_prelu(input_graph_def):
     node.op = 'Identity'
     del node.input[:]
     node.input.append(relu_input_op.name)
+    _create_alpha_node(neg_alpha_op, updated_alpha)
 
     nodes_to_skip[mul_op.name] = True
     nodes_to_skip[relu_neg_input_op.name] = True
@@ -189,4 +189,3 @@ def fuse_prelu_with_fused_conv2d_or_matmul(input_graph_def):
 
   return graph_rewrite_util.cleanup_graph_def(
       input_graph_def, nodes_to_skip, inputs_to_remove)
-

tfjs-converter/python/tensorflowjs/converters/fuse_prelu_test.py

@@ -21,6 +21,10 @@
 import tensorflow.compat.v2 as tf
 from tensorflow.core.protobuf import config_pb2
 from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import variables
+from tensorflow.python.training.tracking import tracking
 
 from tensorflowjs.converters import fuse_depthwise_conv2d
 from tensorflowjs.converters import fuse_prelu
@@ -234,5 +238,27 @@ def execute_model(tensor):
     self.assertNotEqual(conv2d_op, None)
     self.assertEqual(conv2d_op.attr['fused_ops'].list.s, [b'BiasAdd', b'Prelu'])
     self.assertEqual(conv2d_op.attr['num_args'].i, 2)
+
+  def testNonPreluPattern(self):
+    """Test a basic model with functions to make sure functions are inlined."""
+    input_data = constant_op.constant(1., shape=[1])
+    root = tracking.AutoTrackable()
+    root.v1 = variables.Variable(3.)
+    root.v2 = variables.Variable(2.)
+
+    root.f = def_function.function(lambda x: tf.nn.relu(root.v1) + root.v2 * 2.0)
+    to_save = root.f.get_concrete_function(input_data)
+    graph = tf_saved_model_conversion_v2._freeze_saved_model_v2(
+        root.f.get_concrete_function(input_data))
+    graph_def = graph.as_graph_def()
+    graph_def = fuse_prelu.fuse_ops_for_prelu(graph_def)
+    const_op = None
+    for node in graph_def.node:
+      self.assertNotEqual("Prelu", node.op)
+      if node.op == 'Const':
+        const_op = node
+    self.assertNotEqual(const_op, None)
+    self.assertEqual(const_op.attr['value'].tensor.float_val, [2.0])
+
 if __name__ == '__main__':
   tf.test.main()

tfjs-converter/src/executor/graph_model_test.ts

@@ -230,7 +230,7 @@ const INITIALIZER_GRAPHDEF: tensorflow.IGraphDef = {
         use_node_name_sharing: {b: false},
         key_dtype: {type: tensorflow.DataType.DT_STRING},
         container: {s: ''},
-        shared_name: {s: 'tablename'}
+        shared_name: {s: 'dGFibGVuYW1l' /* base64 'tablename' */}
       }
     },
     {
@@ -277,7 +277,7 @@ const HASH_TABLE_MODEL: tensorflow.IGraphDef = {
         use_node_name_sharing: {b: false},
         key_dtype: {type: tensorflow.DataType.DT_STRING},
         container: {s: ''},
-        shared_name: {s: 'tablename'}
+        shared_name: {s: 'dGFibGVuYW1l' /* base64 'tablename' */}
       }
     },
     {

tfjs-core/src/jasmine_util.ts

@@ -217,6 +217,7 @@ export function describeWithFlags(
 
   TEST_ENVS.forEach(testEnv => {
     env().setFlags(testEnv.flags);
+    env().set('IS_TEST', true);
     if (envSatisfiesConstraints(env(), testEnv, constraints)) {
       const testName =
           name + ' ' + testEnv.name + ' ' + JSON.stringify(testEnv.flags || {});

tools/make_rollup_config.ts

@@ -18,7 +18,7 @@
 import * as commonjs_import from '@rollup/plugin-commonjs';
 import resolve from '@rollup/plugin-node-resolve';
 import * as sourcemaps_import from 'rollup-plugin-sourcemaps';
-import * as visualizer from 'rollup-plugin-visualizer';
+import {visualizer} from 'rollup-plugin-visualizer';
 import {terser as terserPlugin} from 'rollup-plugin-terser';
 import {downlevelToEs5Plugin} from 'downlevel_to_es5_plugin/downlevel_to_es5_plugin';