`ValueError: The two structures don't have the same nested structure.`

[DarrenZhang01]

As titled.

empirical.py", line 49, in _jvp
    with forwardprop.ForwardAccumulator(primals, tangents) as acc:
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/eager/forwardprop.py", line 363, in __init__
    self._watch(primals, tangents)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/eager/forwardprop.py", line 414, in _watch
    nest.map_structure(_watch, primals, tangents, expand_composites=True)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/util/nest.py", line 630, in map_structure
    assert_same_structure(structure[0], other, check_types=check_types,
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/util/nest.py", line 401, in assert_same_structure
    raise type(e)("%s\n"
ValueError: The two structures don't have the same nested structure.

First structure: type=tuple str=([(<tf.Tensor 'stateless_random_normal:0' shape=(7, 1024) dtype=float32>, <tf.Tensor 'stateless_random_normal_1:0' shape=(1, 1024) dtype=float32>), (), (<tf.Tensor 'stateless_random_normal_2:0' shape=(1024, 1) dtype=float32>, <tf.Tensor 'stateless_random_normal_3:0' shape=(1, 1) dtype=float32>)],)

Second structure: type=tuple str=([(ndarray<<tf.Tensor 'gradient_tape/Tensordot/transpose/transpose:0' shape=(7, 1024) dtype=float32>>, ndarray<<tf.Tensor 'gradient_tape/mul_1/Mul:0' shape=(1, 1024) dtype=float32>>), (), (ndarray<<tf.Tensor 'gradient_tape/Cast_4/Cast:0' shape=(1024, 1) dtype=float32>>, ndarray<<tf.Tensor 'gradient_tape/mul_5/Mul:0' shape=(1, 1) dtype=float32>>)],)

More specifically: Substructure "type=ndarray str=ndarray<Tensor("gradient_tape/Tensordot/transpose/transpose:0", shape=(7, 1024), dtype=float32)>" is a sequence, while substructure "type=Tensor str=Tensor("stateless_random_normal:0", shape=(7, 1024), dtype=float32)" is not
Entire first structure:
([(., .), (), (., .)],)
Entire second structure:
([(., .), (), (., .)],)

[DarrenZhang01]

Solved in a15151c