Loading saved_model causes "Failed to compile fragment shader" for gather op

I'm using Tensorflow in Python to define and train the model, and then I save it to a saved_model and load it on a website using TensorflowJS. I've made a minimal working example presented below, and have isolated the issue to the gather op. I'm using Windows 7 x64, Python 3.5.2, numpy 1.15.1 (forced upon installing tensorflowjs - had 15.2 before that) and the latest version of tensorflowjs installed using simply pip install tensorflowjs. The browser is Chrome, and it's all served on localhost using Wamp.

Step 1: Python

# coding=utf-8

import tensorflow as tf
import sys
import os, shutil

x = tf.placeholder(tf.int32, shape=[5], name='x')

#pred = tf.transpose(x, name="y") # THIS WORKS
pred = tf.gather(x, 1, axis=0, name="y") # THIS DOES NOT - thus must be related to the gather op

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(pred, feed_dict={x: [0,1,2,3,4]})) # just to check that it works in Python
    tf.saved_model.simple_save(sess, "./export", inputs={"x": x}, outputs={"y": pred})

Step 2: tfjs-converter command

tensorflowjs_converter --input_format=tf_saved_model --output_node_names="y" --saved_model_tags=serve export C:/wamp/www/avatar/web_model

Step 3: Serving the website:

<!-- index.html -->

<!DOCTYPE html>
<html xmlns="http://www.w3.org/1999/xhtml">

<head>
  <meta http-equiv="Content-Type" content="text/html" charset="utf-8" />
  <title>test</title>
  <script src="./ext/jquery-3.3.1.min.js"></script>
  <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script>
</head>

<body>
  <script src="./script.js"></script>
</body>

</html>

// script.js

async function learnLinear() {
  const model = await tf.loadFrozenModel('http://localhost/avatar/web_model/tensorflowjs_model.pb', 'http://localhost/avatar/web_model/weights_manifest.json');

  x = tf.tensor1d([2, 3, 4, 1, 1], "int32");

  model.execute({
    "x": x
  }); // "x: The input data, as an Tensor, or an Array of tf.Tensors if the model has multiple inputs."
}
learnLinear();

The error produced is:

Uncaught (in promise) Error: Failed to compile fragment shader.
    at createFragmentShader (tfjs:2)
    at e.createProgram (tfjs:2)
    at compileProgram (tfjs:2)
    at tfjs:2
    at e.getAndSaveBinary (tfjs:2)
    at e.compileAndRun (tfjs:2)
    at e.gather (tfjs:2)
    at ENV.engine.runKernel.$x (tfjs:2)
    at tfjs:2
    at e.scopedRun (tfjs:2)

And above that in the console there's:

ERROR: 0:157: 'getIndices' : no matching overloaded function found
157          setOutput(getA(int(getIndices(resRC))));

So getIndices is missing, and it's not declared in the tfjs script that's imported, only called...

Any help is much appreciated!

Solution

tf.gather has parameters indices of type tensor. And this is the case both in python and in js

To avoid your error, you have to change the following:

pred = tf.gather(x, tf.constant([1], dtype="int32"), axis=0, name="y")

After building the graph again and converting it to a frozen model you will be able to use it in js.