Create a None Dimension in a Tensorflow dataset

So currently I try to prepare my data to pass it to my Model. I have following Problem. If i understood the concept of datasets right it passes the Input and the y value to my model. If im wrong please correct me. I try to prepare it like this.

dataset = dataset.map(
    lambda  X, modules_sorted, solvingState:(
      tf.one_hot(X, depth = num_concepts*2),
      tf.concat(values=[
        tf.one_hot(modules_sorted, depth = num_concepts,dtype=tf.int32),
        tf.expand_dims(solvingState, axis=-1)
      ],axis=-1
      )
    )
)

My Problem is that the dataset should look like this afterwards.

[[batch_size x max_seq_length x num_features],[batch_size x max_seq_length x num_concepts]]

But my dataset looks like this cause I cant figure a way to expand with a None dimension

    [[batch_size x max_seq_length x num_features],[batch_size x max_seq_length x num_concepts+1]]

Now my question is there a way to expand solvingStates with a None dim so my output can look like I mentioned before?

Thanks for every bit of help.