keras.ops.map can't handle nested structures for TensorFlow backend #20048

apage224 · 2024-07-26T14:59:47Z

Keras: 3.4.1
TensorFlow: 2.17.0

As background, I am looking to leverage both keras.ops.map as well as keras.ops.vectorize_map for custom preprocessing layers. Certain layers require sequential mapping hence I use keras.ops.map. If I pass a nested input, keras.ops.map will fail when using TensorFlow backend.

I believe this line is an issue as it assumes the input is not nested:

def map(f, xs):
    xs = tree.map_structure(convert_to_tensor, xs)

    def get_fn_output_signature(x):
        out = f(x)
        return tree.map_structure(tf.TensorSpec.from_tensor, out)

    fn_output_signature = get_fn_output_signature(xs[0])
    return tf.map_fn(f, xs, fn_output_signature=fn_output_signature)

From what I can tell, it is trying to determine the output signature (which might not match the input) by feeding the function a single element (e.g. xs[0]) which won't work on nested inputs. I was able to fix it by updating the function as follows (note: I've done only limited testing).

def map(f, xs):
    xs = tree.map_structure(convert_to_tensor, xs)

    def get_fn_output_signature(x):
        out = f(x)
        return tree.map_structure(tf.TensorSpec.from_tensor, out)
    # Grab single element unpacking and repacking single element        
    x = tf.nest.pack_sequence_as(xs, [x[0] for x in tf.nest.flatten(xs)])
    fn_output_signature = get_fn_output_signature(x)  
    return tf.map_fn(f, xs, fn_output_signature=fn_output_signature)

Test case:

import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import keras
import tensorflow as tf

def my_fn(inputs):    
    outputs = dict(inputs)
    outputs['x'] = inputs['x'][:, 0]
    outputs['y'] = inputs['y']   1
    return outputs

xs = {
    'x': tf.convert_to_tensor(np.random.rand(4, 100, 3), dtype=tf.float32),
    'y': tf.convert_to_tensor(np.random.randint(0, 10, size=(4, 1)), dtype=tf.int32)
}

Calling keras.ops.map:

ys = keras.ops.map(my_fn, xs)

produces error:

    225     out = f(x)
    226     return tree.map_structure(tf.TensorSpec.from_tensor, out)
--> 228 fn_output_signature = get_fn_output_signature(xs[0])
    229 return tf.map_fn(f, xs, fn_output_signature=fn_output_signature)

KeyError: 0"

Calling custom map:

ys = map(my_fn, xs)
print(ys['x'].shape)

produces correct result:

    (4, 100)

The text was updated successfully, but these errors were encountered:

mehtamansi29 · 2024-07-30T08:52:12Z

Hi @apage224 -

Thanks for raising this issue. We have raised an internal fix for this issue. Soon it will be reflected in the documentation as well. This fix is working fine for nested structure(i.e dict,list,tuple).

Example:

import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import keras
import tensorflow as tf
import numpy as np

def my_fn(inputs): 
    outputs = [x**2 for x in inputs]
    return outputs

xs = [tf.convert_to_tensor(np.random.rand(4, 100, 3), dtype=tf.float32),tf.convert_to_tensor(np.random.randint(0, 10, size=(4, 1)), dtype=tf.int32)]

ys = keras.ops.map(my_fn, xs)
print(ys)

github-actions bot assigned mehtamansi29 Jul 26, 2024

mehtamansi29 added type:Bug backend:tensorflow labels Jul 29, 2024

mehtamansi29 linked a pull request Jul 30, 2024 that will close this issue

Handled nested structures for TensorFlow backend for keras.ops.map. #20064

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

keras.ops.map can't handle nested structures for TensorFlow backend #20048

keras.ops.map can't handle nested structures for TensorFlow backend #20048

apage224 commented Jul 26, 2024

mehtamansi29 commented Jul 30, 2024 •

edited

Loading

keras.ops.map can't handle nested structures for TensorFlow backend #20048

keras.ops.map can't handle nested structures for TensorFlow backend #20048

Comments

apage224 commented Jul 26, 2024

mehtamansi29 commented Jul 30, 2024 • edited Loading

mehtamansi29 commented Jul 30, 2024 •

edited

Loading