Unsupported ops: TensorListStack, TensorListReserve, TensorListFromTensor

[alvinsunyixiao]

TensorFlow.js

2.0.1

Browser

Chrome 83.0.4103.116

Problem Description

I am trying to implement a real-time interactive high-dimensional ODE simulation tool with tfjs. My approach is by defining the complicated computation in Python and generate TF Model as computation graphs, and then trigger the simulation step in JavaScript by performing inference on the saved graph (TF Saved Model). I run into this error when running tensorflowjs_converter on a custom model. I believe tensorflow_probability.math.ode.DormandPrince.solve in my model introduced those unsupported ops.

Feature Request

I am hoping if TensorListFromTensor can be supported. However, I am not entirely sure if the root cause of the problem is TensorListFromTensor. According to this, it seems like the other 2 ops (TensorListStack and TensorListReserve) are actually supported by the latest release, but the error message still indicates that they are unsupported

Code to reproduce the Problem

Below is an ODE for simulating a toy version microfluidic equation. The code will save the computation graph to <saved_model_path>.

import tensorflow as tf
import tensorflow_probability as tfp

class SimModel(tf.Module):
    @tf.function(input_signature=(tf.TensorSpec([None, None], tf.float32, name='concentration_sx'), 
                                  tf.TensorSpec([None], tf.float32, name='alpha_s'),
                                  tf.TensorSpec([], tf.float32, name='dx'),
                                  tf.TensorSpec([], tf.float32, name='t_init'),
                                  tf.TensorSpec([], tf.float32, name='t_final')))
    def __call__(self, concentration_sx, alpha_s, dx, t_init, t_final):
        ode_fn = lambda t, y: self.calc_flux(y, alpha_s, dx)
        t = tf.linspace(t_init, t_final, 10)
        solver = tfp.math.ode.DormandPrince()
        results = solver.solve(ode_fn, t_init, concentration_sx, t)
        return results.times, results.states

    def calc_flux(self, concentration_sx, alpha_s, dx):
        c_left_sx = concentration_sx[:, :-1]
        c_right_sx = concentration_sx[:, 1:]
        c_left_sum_1x = tf.reduce_sum(c_left_sx, axis=0, keepdims=True)
        c_right_sum_1x = tf.reduce_sum(c_right_sx, axis=0, keepdims=True)
        adv_flux_sx = alpha_s[:, tf.newaxis] * 0.5 * (c_left_sx / c_left_sum_1x + c_right_sx / c_right_sum_1x)
        v_max_1x = tf.reduce_max(alpha_s) * 0.5 * (1 / c_left_sum_1x + 1 / c_right_sum_1x)
        num_diff_sx = 0.5 * tf.abs(v_max_1x) * (c_right_sx - c_left_sx)
        flux_sx = adv_flux_sx - num_diff_sx
        rhs_sx = (flux_sx[:, :-1] - flux_sx[:, 1:]) / dx
        return tf.pad(rhs_sx, [[0, 0], [1, 1]])

model = SimModel()
tf.saved_model.save(model, '<saved_model_path>')

And then calling the following conversion script triggers this error.

tensorflowjs_converter --control_flow_v2=True --input_format=tf_saved_model --saved_model_tags=serve --signature_name=serving_default --strip_debug_ops=True --weight_shard_size_bytes=4194304 <saved_model_path> <output_model_path>

[pyu10055]

@alvinsunyixiao Those tensorlist ops has been added, you might need to wait for next release, hopeful it can be done this week.

[rthadur]

Closing this , feel free to reopen if you see any issue after the release happens.