Support Complex64 weights conversion and loading

.gitignore

@@ -28,6 +28,7 @@ tfjs-layers/integration_tests/tfjs2keras/test-data/
 tfjs-layers/integration/typescript/yarn.lock
 e2e/integration_tests/create_save_predict_data
 e2e/integration_tests/convert_predict_data
+tfjs-converter/python/tensorflowjs.egg-info
 
 # Ignore the src, binding, scripts of tfjs-node-gpu since it is copied over when
 # building.

e2e/integration_tests/constants.ts

@@ -34,7 +34,8 @@ export const LAYERS_MODELS = [
 
 export const GRAPH_MODELS = [
   'saved_model_v1', 'saved_model_v2', 'saved_model_v2_with_control_flow',
-  'saved_model_with_conv2d', 'saved_model_with_prelu'
+  'saved_model_with_conv2d', 'saved_model_with_prelu',
+  'saved_model_v2_complex64', 'saved_model_v2_with_control_flow_v2'
 ];
 
 /** Karma server directory serving local files. */

e2e/integration_tests/convert_predict.py

@@ -50,7 +50,7 @@
 curr_dir = os.path.dirname(os.path.realpath(__file__))
 _tmp_dir = os.path.join(curr_dir, 'convert_predict_data')
 
-def _save_and_convert_model(model_fn, model_path):
+def _save_and_convert_model(model_fn, model_path, control_flow_v2=False):
   """Benchmark a model's fit() and predict() calls; serialize the model.
 
   Args:
@@ -94,13 +94,17 @@ def _save_and_convert_model(model_fn, model_path):
   artifacts_dir = os.path.join(_tmp_dir, model_path)
 
   # Convert and store model to file.
-  subprocess.check_output([
+  args = [
       'tensorflowjs_converter',
       '--input_format', 'tf_saved_model',
       '--output_format', 'tfjs_graph_model',
       '--signature_name', 'serving_default',
-      '--saved_model_tags', 'serve',
-      tmp_saved_model_dir, artifacts_dir])
+      '--saved_model_tags', 'serve'];
+  if control_flow_v2:
+    args = args + ['--control_flow_v2', 'True']
+
+  print(args)
+  subprocess.check_output(args +[tmp_saved_model_dir, artifacts_dir])
 
 def _create_saved_model_v1(save_dir):
   """Create a TensorFlow V1 SavedModel for testing.
@@ -257,6 +261,58 @@ def _create_saved_model_with_prelu(save_dir):
                          "shape": result.shape,
                          "dtype": "float32"}}}
 
+def _create_saved_model_v2_complex64(save_dir):
+  """Test a TF V2 model with complex dtype.
+
+  Args:
+    save_dir: directory name of where the saved model will be stored.
+  """
+  input_data = constant_op.constant(1., shape=[1])
+  root = tracking.AutoTrackable()
+  root.v1 = variables.Variable(3 + 1j, dtype=tf.complex64)
+  root.f = def_function.function(lambda x: tf.complex(x, x) + root.v1)
+  to_save = root.f.get_concrete_function(input_data)
+
+  save(root, save_dir, to_save)
+  return {
+      "async": False,
+      "inputs": {
+          "x": {"value": [1], "shape": [1], "dtype": 'float32'}},
+      "outputs": {
+          "Identity:0": {"value": [4, 2], "shape": [1], "dtype": "complex64"}}}
+
+def _create_saved_model_v2_with_control_flow_v2(save_dir):
+  """Test a TF V2 model with complex dtype.
+
+  Args:
+    save_dir: directory name of where the saved model will be stored.
+  """
+  class CustomModule(tf.Module):
+
+      def __init__(self):
+          super(CustomModule, self).__init__()
+
+      @tf.function(input_signature=[
+          tf.TensorSpec([], tf.float32), tf.TensorSpec([], tf.float32)])
+      def control_flow(self, x, y):
+          while x < y:
+                  x = x + 2
+          return x
+
+
+  module = CustomModule()
+  print(module.control_flow(1, 2))
+  tf.saved_model.save(module, save_dir,
+                      signatures=module.control_flow)
+
+  return {
+      "async": False,
+      "inputs": {
+          "x": {"value": [-1.], "shape": [], "dtype": 'float32'},
+          "y": {"value": [2.], "shape": [], "dtype": 'float32'}},
+      "outputs": {
+          "Identity:0": {"value": [3.], "shape": [], "dtype": "float32"}}}
+
 def main():
   # Create the directory to store model and data.
   if os.path.exists(_tmp_dir) and os.path.isdir(_tmp_dir):
@@ -265,8 +321,12 @@ def main():
 
   _save_and_convert_model(_create_saved_model_v1, 'saved_model_v1')
   _save_and_convert_model(_create_saved_model_v2, 'saved_model_v2')
+  _save_and_convert_model(_create_saved_model_v2_complex64,
+                          'saved_model_v2_complex64')
   _save_and_convert_model(_create_saved_model_v2_with_control_flow,
       'saved_model_v2_with_control_flow')
+  _save_and_convert_model(_create_saved_model_v2_with_control_flow_v2,
+      'saved_model_v2_with_control_flow_v2', control_flow_v2=True)
   _save_and_convert_model(_create_saved_model_with_conv2d,
       'saved_model_with_conv2d')
   _save_and_convert_model(_create_saved_model_with_prelu,

e2e/integration_tests/convert_predict.ts

@@ -60,13 +60,13 @@ describe(`${REGRESSION} convert_predict`, () => {
 
       BACKENDS.forEach(backend => {
         it(`with ${backend}.`, async () => {
+          await tfc.setBackend(backend);
+
           const $model = await tfconverter.loadGraphModel(
               `${KARMA_SERVER}/${DATA_URL}/${model}/model.json`);
 
           const xs = createInputTensors(inputsData, inputsShapes);
 
-          await tfc.setBackend(backend);
-
           const result = await $model.executeAsync(xs);
 
           const ys =

e2e/integration_tests/requirements-dev.txt

@@ -1,3 +1,2 @@
 keras==2.3.1
-tensorflowjs>=1.2.10.1
-tf-nightly>=2.1.0.dev20191007
+-r ../../tfjs-converter/python/requirements.txt

e2e/integration_tests/requirements-stable.txt

@@ -1,3 +1,7 @@
 keras==2.3.1
-tensorflow==1.15.0
-tensorflowjs==1.7.4
+h5py>=2.8.0
+numpy>=1.16.4
+six>=1.12.0
+tensorflow-cpu==2.1.0
+tensorflow-hub==0.7.0
+PyInquirer==1.0.3

e2e/scripts/setup-py-env.sh

@@ -62,3 +62,6 @@ if [[ "${DEV_VERSION}" == "stable" ]]; then
 else
   pip3 install -r requirements-dev.txt
 fi
+
+echo "Loading tensorflowjs pip from source ...."
+pip3 install -e ../../tfjs-converter/python

tfjs-converter/python/build-pip-package.sh

@@ -32,6 +32,7 @@
 #   --test-nightly: Test the pip packages by installing it (inside virtualenv)
 #           and running test_pip_package.py and test_pip_nightly_package.py
 #           against the install.
+#   --build: Create the pip packages.
 #   --upload:         Upload the py2 and py3 wheels to prod PyPI.
 #   --upload-to-test: Upload the py2 and py3 wheels to test PyPI, mutually
 #                     exclusive with --upload.
@@ -62,7 +63,7 @@ set -e
 function print_usage() {
   echo "Usage:"
   echo "  build-pip-packages.sh \\"
-  echo "      [--test] [--test-nightly] [--upload] [--upload-to-test] [--confirm-upload] <OUTPUT_DIR>"
+  echo "      [--test] [--test-nightly] [--build] [--upload] [--upload-to-test] [--confirm-upload] <OUTPUT_DIR>"
   echo
 }
 
@@ -78,6 +79,7 @@ RUN_TEST_NIGHTLY=0
 UPLOAD_TO_PROD_PYPI=0
 UPLOAD_TO_TEST_PYPI=0
 CONFIRM_UPLOAD=0
+BUILD=0
 DEST_DIR=""
 while true; do
   if [[ "$1" == "--test" ]]; then
@@ -90,6 +92,8 @@ while true; do
     UPLOAD_TO_TEST_PYPI=1
   elif [[ "$1" == "--confirm-upload" ]]; then
     CONFIRM_UPLOAD=1
+  elif [[ "$1" == "--build" ]]; then
+    BUILD=1
   elif [[ "$1" != --* ]]; then
     DEST_DIR="$1"
   else

tfjs-converter/python/requirements.txt

@@ -1,6 +1,6 @@
 h5py>=2.8.0
 numpy>=1.16.4
 six>=1.12.0
-tensorflow-cpu>=2.1.0<3
+tensorflow-cpu==2.1.0
 tensorflow-hub==0.7.0
 PyInquirer==1.0.3

tfjs-converter/python/tensorflowjs/read_weights.py

@@ -24,7 +24,8 @@
 import numpy as np
 from tensorflowjs import quantization
 
-_INPUT_DTYPES = [np.float32, np.int32, np.uint8, np.uint16, np.object]
+_INPUT_DTYPES = [np.float32, np.int32, np.complex64,
+                 np.uint8, np.uint16, np.object]
 
 # Number of bytes used to encode the length of a string in a string tensor.
 STRING_LENGTH_NUM_BYTES = 4

tfjs-converter/python/tensorflowjs/read_weights_test.py

@@ -45,6 +45,9 @@ def testReadOneGroup(self):
         [{
             'name': 'weight1',
             'data': np.array([1, 2, 3], 'float32')
+        }, {
+            'name': 'weight2',
+            'data': np.array([1 + 1j, 2 + 2j, 3 + 3j])
         }]
     ]
 
@@ -53,11 +56,13 @@ def testReadOneGroup(self):
     # Read the weights using `read_weights`.
     read_output = read_weights.read_weights(manifest, self._tmp_dir)
     self.assertEqual(1, len(read_output))
-    self.assertEqual(1, len(read_output[0]))
+    self.assertEqual(2, len(read_output[0]))
     self.assertEqual('weight1', read_output[0][0]['name'])
     self.assertTrue(
         np.allclose(groups[0][0]['data'], read_output[0][0]['data']))
-
+    self.assertEqual('weight2', read_output[0][1]['name'])
+    self.assertTrue(
+        np.allclose(groups[0][1]['data'], read_output[0][1]['data']))
   def testReadOneGroupString(self):
     groups = [
         [{

tfjs-converter/python/tensorflowjs/write_weights.py

@@ -24,10 +24,12 @@
 from tensorflowjs import quantization
 from tensorflowjs import read_weights
 
-_OUTPUT_DTYPES = [np.float32, np.int32, np.uint8, np.uint16, np.bool, np.object]
+_OUTPUT_DTYPES = [np.float32, np.int32, np.complex64,
+                  np.uint8, np.uint16, np.bool, np.object]
 _AUTO_DTYPE_CONVERSION = {
     np.dtype(np.float64): np.float32,
-    np.dtype(np.int64): np.int32}
+    np.dtype(np.int64): np.int32,
+    np.dtype(np.complex128): np.complex64}
 
 def write_weights(
     weight_groups, write_dir, shard_size_bytes=1024 * 1024 * 4,

tfjs-converter/python/tensorflowjs/write_weights_test.py

@@ -278,6 +278,37 @@ def test_1_group_1_weight_string_sharded(self):
     string = weight_bytes[4:14].decode('utf-8')
     self.assertEqual(string, u'helloworld')
 
+  def test_1_group_1_weight_complex(self):
+    groups = [
+        [{
+            'name': 'weight1',
+            'data': np.array([1 + 1j, 2 + 2j, 3 + 3j], 'complex')
+        }]
+    ]
+
+    manifest = write_weights.write_weights(
+        groups, TMP_DIR, shard_size_bytes=6 * 4)
+
+    self.assertTrue(
+        os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')),
+        'weights_manifest.json does not exist')
+
+    self.assertEqual(
+        manifest,
+        [{
+            'paths': ['group1-shard1of1.bin'],
+            'weights': [{
+                'name': 'weight1',
+                'shape': [3],
+                'dtype': 'complex64'
+            }]
+        }])
+
+    weights_path = os.path.join(TMP_DIR, 'group1-shard1of1.bin')
+    weight1 = np.fromfile(weights_path, 'complex64')
+    np.testing.assert_array_equal(
+        weight1, np.array([1 + 1j, 2 + 2j, 3 + 3j], 'complex64'))
+
   def test_1_group_3_weights_packed_multi_dtype(self):
     # Each string tensor uses different encoding.
     groups = [
@@ -654,17 +685,6 @@ def test_bad_weights_entry_throws_no_data(self):
     with self.assertRaises(Exception):
       write_weights.write_weights(groups, TMP_DIR)
 
-  def test_bad_numpy_array_dtype_throws(self):
-    groups = [
-        [{
-            'name': 'weight1',
-            'data': np.array([1, 2, 3], 'complex')
-        }]
-    ]
-
-    with self.assertRaises(Exception):
-      write_weights.write_weights(groups, TMP_DIR)
-
   def test_duplicate_weight_name_throws(self):
     groups = [
         [{

tfjs-core/src/io/io_utils.ts

@@ -15,6 +15,8 @@
  * =============================================================================
  */
 
+import {complex} from '../ops/complex_ops';
+
 import {tensor} from '../ops/tensor_ops';
 import {NamedTensor, NamedTensorMap} from '../tensor_types';
 import {TypedArray} from '../types';
@@ -57,7 +59,7 @@ export async function encodeWeights(
     const name = names[i];
     const t = Array.isArray(tensors) ? tensors[i].tensor : tensors[name];
     if (t.dtype !== 'float32' && t.dtype !== 'int32' && t.dtype !== 'bool' &&
-        t.dtype !== 'string') {
+        t.dtype !== 'string' && t.dtype !== 'complex64') {
       throw new Error(`Unsupported dtype in weight '${name}': ${t.dtype}`);
     }
     const spec: WeightsManifestEntry = {name, shape: t.shape, dtype: t.dtype};
@@ -171,13 +173,25 @@ export function decodeWeights(
         values = new Int32Array(byteBuffer);
       } else if (dtype === 'bool') {
         values = new Uint8Array(byteBuffer);
+      } else if (dtype === 'complex64') {
+        values = new Float32Array(byteBuffer);
+        const real = new Float32Array(values.length / 2);
+        const image = new Float32Array(values.length / 2);
+        for (let i = 0; i < real.length; i++) {
+          real[i] = values[i * 2];
+          image[i] = values[i * 2 + 1];
+        }
+        const realTensor = tensor(real, shape, 'float32');
+        const imageTensor = tensor(image, shape, 'float32');
+        out[name] = complex(realTensor, imageTensor);
       } else {
         throw new Error(`Unsupported dtype in weight '${name}': ${dtype}`);
       }
       offset += size * dtypeFactor;
     }
-
-    out[name] = tensor(values, shape, dtype);
+    if (dtype !== 'complex64') {
+      out[name] = tensor(values, shape, dtype);
+    }
   }
   return out;
 }