Tensorflow py_function setting return value shape

user_89904 0
Viewed 783

I am tring to set the shape of an numpy array return from a tf.py_function call. Function returns a tuple of array and its shape, so I can set its shape inside tf.function. However I can not use the returned shape array inside a set_shape call. I tried casting it to int32 or accessing it via indices etc. but all results in a TypeError: Dimension value must be integer or None or have an __index__ method, got value '<tf.Tensor 'strided_slice:0' shape=<unknown> dtype=int32>' with type '<class 'tensorflow.python.framework.ops.Tensor'>' error.

@tf.function
def samples(img, mask) -> tuple:
    xs, xs_shape = tf.py_function(process,[128, 128], [tf.float32, tf.int32])
    xs.set_shape(tf.TensorShape(xs_shape))
    return xs

ds = ds.map(samples)
tensorflow2.0 keras tensorflow python
1 Answers

This is a contrived example. I assume this could be your requirement.

def f(x):
    print("print ", x)
    # Include Logic using 'x' which is a tensor.
    # Sample  return values
    return [tf.convert_to_tensor([1,2]), tf.convert_to_tensor([1,2]).shape]


@tf.function
def samples(x) -> tuple:
    xs, xs_shape = f(x)
    xs.set_shape(xs_shape)
    print( xs_shape)
    return xs

dataset = tf.data.Dataset.range(1, 6)
dataset = dataset.map(  samples )
user_12349090
Related