Tensorflow tf.dataset won't iterate with multiple inputs of different sizes "Shapes of all inputs must match"

Question

I am trying to make a tensorflow model with two different inputs, one will have shape [9,10], the other will just have shape [8].

I am furthermore trying to use tf.dataset to iterate over my inputs. However, whenever I try to do so it fails with the following error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Shapes of all inputs must match: values[0].shape = [9,10] != values[1].shape = [8]
     [[{{node packed}}]] [Op:IteratorGetNext]

But surely it is possible to have differently sized inputs into different branches! This is exactly the case in the example in tensorflow's functional API guide, however they do not use tf.dataset so I can't simply follow their example.

To give a little more specifics into the problem I am trying to solve and why I am using the tf.dataset api:

I am doing a time-series problem over multiple sites where my inputs are of two types: those that vary with time, and those that do not but do vary by site. For the time being, I'm just trying to estimate the next time step.

First, I get my dynamic covariates and targets in a sliding window using the timeseries_dataset_from_array util.

    train_ds = tf.keras.preprocessing.timeseries_dataset_from_array(
        input_data, targets, sequence_length=window_size, batch_size=256)

This works perfectly and I can train models using this dataset as is.

However, I want to also use the static covariates from the specific site that the time series data is coming from. The site id is included in the window input data in its own column, though it gets removed before training. Thus, what I am trying to do is grab the static covariates for the site and attach it as a separate input to my dataset.

    train_ds = train_ds.map(lambda x, y: (tf.py_function(attach_static_covariates, [x, idindex, colnames], [tf.float64, tf.float64]), y))

    train_ds = train_ds.map(lambda x, y: ({'dynamic': x[0], 'static': x[1]}, y))

The code for the attach_static_covariates method is:

def attach_static_covariates(x, idindex, colnames):
    id = x[0, idindex].numpy()
    static_cov = static_df.iloc[int(id)]
    #This just filters out the id column, now that it has served its purpose
    x = tf.gather(x, [i for i in range(len(colnames)) if i != idindex])

    return (x, static_cov)

I've confirmed that my code can run and train on multiple inputs provided by the above method provided they are the same size (e.g. if I return (x, x) I can run my model on two copies of the dynamic covariates inputted into two different branches of my model). The problem is not due to a mismatch or a bad model definition because I get the same error from the following code:

    for feature_batch, label_batch in train_ds.take(1):
        print(feature_batch)

I've looked everywhere on google and the tensorflow git and I can't find anyone else with this problem, and yet it surely MUST be possible to have differently shaped inputs using tf.dataset! I can't imagine that such an incredibly common use case would be completely unsupported. However, I can't find any examples online where someone has multiple inputs of different shapes and uses tf.dataset api. Any help or links to such examples would be greatly appreciated.

Colab notebook to illustrate issue: https://colab.research.google.com/drive/1EnaJoUULl-fyAwlcG_5tcWfsRFVCOMtv#scrollTo=PHvsIOx6-Uux

Answer 1

I have discovered a workaround: embed the py_function, which can't return a dict, inside another function that is mapped to return a dict, thereby sidestepping the issue of the dataset ever iterating over nondicted double inputs (which is where the bug happens because it interprets them as being batched together and thus demands they have the same dims).

E.g.

    def dictionarize(x):

        x, static_cov = tf.py_function(attach_static_covariates, [x, idindex, colnames], [tf.float64, tf.float64])

        return {'dynamic': x, 'static': static_cov}

    train_ds = train_ds.map(lambda x, y: (dictionarize(x), y))