rank error when trying to use an RNN on Shakespeare text - tensorflow

Question

I am trying to create an RNN that can generate Shakespeare like text. However, I keep on getting an error about the rank of my input matrices/datasets. I have tried squeezing the matrices, and changing the dimensions etc., but I cannot figure out why it is not working.
Here is my code:
# Get Shakespeare text
import tensorflow as tf
shakespeare_url="https://homl.info/shakespeare"
filepath=tf.keras.utils.get_file('shakespeare.txt',shakespeare_url)
with open(filepath)as f:
shakespeare_text=f.read()
#Tokenizing shakespeare text
tokenizer=tf.keras.preprocessing.text.Tokenizer(char_level=True)
tokenizer.fit_on_texts(shakespeare_text)
max_id=len(tokenizer.word_index)
import numpy as np
encoded=np.array(tokenizer.texts_to_sequences(shakespeare_text))-1
dataset_size=tokenizer.document_count
train_size=dataset_size*90//100
dataset=tf.data.Dataset.from_tensor_slices(encoded[:train_size])
#Put dataset object into windows and batches
n_steps=100
window_length=n_steps+1
dataset=dataset.window(window_length,shift=1, drop_remainder=True)
dataset=dataset.flat_map(lambda window: window.batch(window_length))
batch_size=32
dataset=dataset.shuffle(10000).batch(batch_size)
dataset=dataset.map(lambda window: (window[:,:-1],window[:,1:]))
dataset=dataset.map(lambda X_batch, Y_batch: (tf.one_hot(X_batch, depth=max_id), Y_batch))
new_dataset=dataset.map(lambda X_batch, Y_batch: (tf.squeeze(X_batch),tf.squeeze(Y_batch)))
dataset=dataset.prefetch(1)
new_dataset=new_dataset.prefetch(1)
#make model
model=tf.keras.models.Sequential([
#tf.keras.layers.Flatten(input_shape=[batch_size,None,max_id]),
tf.keras.layers.GRU(128, return_sequences=True, stateful=True, dropout=0.2,recurrent_dropout=0.2,
batch_input_shape=[batch_size,None,max_id],name='gru_1'),
tf.keras.layers.GRU(128, return_sequences=True, stateful=True, dropout=0.2,recurrent_dropout=0.2,name='gru_2'),
#batch_input_shape=[batch_size,100,max_id]),
tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(max_id,activation='softmax',name='time_layer'))
])
class ResetStatesCallback(tf.keras.callbacks.Callback):
def on_epoch_begin(self,epoch,logs):
self.model.reset_states()
model.compile(loss='sparse_categorical_crossentropy', optimizer="adam")
model.fit(new_dataset, epochs=50, callbacks=[ResetStatesCallback()])
I have tried passing dataset and new _dataset, both without success.
The Error is below:
Epoch 1/50
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-329-28b67720fd03> in <module>
1 model.compile(loss='sparse_categorical_crossentropy', optimizer="adam")
----> 2 model.fit(new_dataset, epochs=50, callbacks=[ResetStatesCallback()])
~/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py in _method_wrapper(self, *args, **kwargs)
106 def _method_wrapper(self, *args, **kwargs):
107 if not self._in_multi_worker_mode(): # pylint: disable=protected-access
--> 108 return method(self, *args, **kwargs)
109
110 # Running inside `run_distribute_coordinator` already.
~/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)
1096 batch_size=batch_size):
1097 callbacks.on_train_batch_begin(step)
-> 1098 tmp_logs = train_function(iterator)
1099 if data_handler.should_sync:
1100 context.async_wait()
~/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
778 else:
779 compiler = "nonXla"
--> 780 result = self._call(*args, **kwds)
781
782 new_tracing_count = self._get_tracing_count()
~/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
821 # This is the first call of __call__, so we have to initialize.
822 initializers = []
--> 823 self._initialize(args, kwds, add_initializers_to=initializers)
824 finally:
825 # At this point we know that the initialization is complete (or less
~/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
695 self._concrete_stateful_fn = (
696 self._stateful_fn._get_concrete_function_internal_garbage_collected( # pylint: disable=protected-access
--> 697 *args, **kwds))
698
699 def invalid_creator_scope(*unused_args, **unused_kwds):
~/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
2853 args, kwargs = None, None
2854 with self._lock:
-> 2855 graph_function, _, _ = self._maybe_define_function(args, kwargs)
2856 return graph_function
2857
~/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
3211
3212 self._function_cache.missed.add(call_context_key)
-> 3213 graph_function = self._create_graph_function(args, kwargs)
3214 self._function_cache.primary[cache_key] = graph_function
3215 return graph_function, args, kwargs
~/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
3073 arg_names=arg_names,
3074 override_flat_arg_shapes=override_flat_arg_shapes,
-> 3075 capture_by_value=self._capture_by_value),
3076 self._function_attributes,
3077 function_spec=self.function_spec,
~/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
984 _, original_func = tf_decorator.unwrap(python_func)
985
--> 986 func_outputs = python_func(*func_args, **func_kwargs)
987
988 # invariant: `func_outputs` contains only Tensors, CompositeTensors,
~/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
598 # __wrapped__ allows AutoGraph to swap in a converted function. We give
599 # the function a weak reference to itself to avoid a reference cycle.
--> 600 return weak_wrapped_fn().__wrapped__(*args, **kwds)
601 weak_wrapped_fn = weakref.ref(wrapped_fn)
602
~/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
971 except Exception as e: # pylint:disable=broad-except
972 if hasattr(e, "ag_error_metadata"):
--> 973 raise e.ag_error_metadata.to_exception(e)
974 else:
975 raise
ValueError: in user code:
/Users/dariyankhan/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py:806 train_function *
return step_function(self, iterator)
/Users/dariyankhan/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py:796 step_function **
outputs = model.distribute_strategy.run(run_step, args=(data,))
/Users/dariyankhan/anaconda3/lib/python3.7/site-packages/tensorflow/python/distribute/distribute_lib.py:1211 run
return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
/Users/dariyankhan/anaconda3/lib/python3.7/site-packages/tensorflow/python/distribute/distribute_lib.py:2585 call_for_each_replica
return self._call_for_each_replica(fn, args, kwargs)
/Users/dariyankhan/anaconda3/lib/python3.7/site-packages/tensorflow/python/distribute/distribute_lib.py:2945 _call_for_each_replica
return fn(*args, **kwargs)
/Users/dariyankhan/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py:789 run_step **
outputs = model.train_step(data)
/Users/dariyankhan/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py:747 train_step
y_pred = self(x, training=True)
/Users/dariyankhan/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:976 __call__
self.name)
/Users/dariyankhan/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/input_spec.py:168 assert_input_compatibility
layer_name + ' is incompatible with the layer: '
ValueError: Input 0 of layer sequential_34 is incompatible with the layer: its rank is undefined, but the layer requires a defined rank.