Dynamic while loop Tensorflow

Question

I have this numpy function, I simplified it into something like this:

def _func(new_inputs, X):
    """ Basically any operation. Below is just an example """
    outputs = new_inputs + X
    new_inputs.pop(0)
    return outputs, new_inputs

new_inputs = []
flag = True
while flag:
    outputs, new_inputs = _func(new_inputs, X)
    X = np.concatenate([X, outputs], axis=0)
    if not new_inputs:
        flag = False

I want to convert it into equivalent TF functions that can be supported in AutoGraph. I'm using TF 2.4.1 but I need to run it under graph computation because I need to run it under Beam context.

This is my attempt:

def function_to_solve():
    new_inputs = tf.constant([])
    flag = tf.constant(True)
    outputs = tf.constant([])
    X = tf.TensorArray(dtype=tf.float32, infer_shape=False, size=1,
                        dynamic_size=True)

    def _func(new_inputs, X):
        return outputs, new_inputs


    def loop_func(flag, new_inputs, outputs, X, i):
        outputs, new_inputs = _func(new_inputs, X)
        X.write(i, outputs)
        if i==2:
            flag = False
        return flag, new_inputs, outputs, X, i+1


    def condition(flag, *args):
        return flag


    _, _, _, X, _ = tf.while_loop(condition, loop_func, [flag, new_inputs, outputs, X, 0])
    X.concat()
    return X

Then call it under Beam context:

import pprint
import tempfile

import tensorflow as tf
import tensorflow_transform as tft
import tensorflow_transform.beam.impl as tft_beam
from tensorflow_transform.tf_metadata import dataset_metadata
from tensorflow_transform.tf_metadata import dataset_schema


raw_data = [
      {'x': [1, 2, 0, 0, 0, 0.3, 10, 1, 1.4]},
 ]

raw_data_metadata = dataset_metadata.DatasetMetadata(
    dataset_schema.from_feature_spec({
        'x': tf.io.FixedLenFeature(shape=(9,), dtype=tf.float32),
    }))

def preprocessing_fn(inputs):
    outputs = function_to_solve()
    return {
        'outputs ': outputs
    }

with tft_beam.Context(temp_dir=tempfile.mkdtemp()):
  transformed_dataset, transform_fn = (
      (raw_data, raw_data_metadata) | tft_beam.AnalyzeAndTransformDataset(
          preprocessing_fn))

transformed_data, transformed_metadata = transformed_dataset

print('\nRaw data:\n{}\n'.format(pprint.pformat(raw_data)))
print('Transformed data:\n{}'.format(pprint.pformat(transformed_data)))

I got the following error:

AttributeError: 'TensorArray' object has no attribute 'get_shape'

Can anyone please help? Thanks!

score 1 · Answer 1 · answered May 11 '21 at 12:40

Every tf.TensorArray object can be transformed into a tf.Tensor using the ta.stack() function.

tensor = ta.stack()

So wherever the error is happening, you must convert the TensorArray object to a regular Tensor. The get_shape function is currently called inside another TF function.

Dynamic while loop Tensorflow

1 Answers1