TypeError when using dynamic_unroll

[kinalmehta]

Error

TypeError: scan carry output and input must have same type structure, got PyTreeDef((CustomNode(namedtuple[<class 'haiku._src.recurrent.LSTMState'>], [*, *]), CustomNode(namedtuple[<class 'haiku._src.stateful.InternalState'>], [None, CustomNode(<class 'collections.defaultdict'>[(<class 'dict'>, ())], []), None]))) and PyTreeDef((CustomNode(namedtuple[<class 'tensorflow.python.saved_model.nested_structure_coder.LSTMState'>], [*, *]), CustomNode(namedtuple[<class 'haiku._src.stateful.InternalState'>], [None, CustomNode(<class 'collections.defaultdict'>[(<class 'dict'>, ())], []), None]))).

This does not occur when using static_unroll even when the datatypes are the same as in the above case. This seems to be some issue with the scan function used inside dynamic_unroll.