How to preserve dict keys in tf.data.Dataset.from_generator?

Question

I have data with non rectangular shape like this:

samples_train = {\'data\': [np.array([[1,1]]), np.array([[1,1],[2,2]]