`assertEqual` Does not Support Trax `ndarray` in the Updated Neural Tangents Test Case

[DarrenZhang01]

    self.assertEqual(x_dict[field.name], y_dict[field.name])
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/unittest/case.py", line 912, in assertEqual
    assertion_func(first, second, msg=msg)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/unittest/case.py", line 902, in _baseAssertEqual
    if not first == second:
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/trax/tf_numpy/numpy_impl/arrays.py", line 172, in __bool__
    return self.__nonzero__()
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/trax/tf_numpy/numpy_impl/arrays.py", line 169, in __nonzero__
    return bool(self.data)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 884, in __bool__
    return bool(self._numpy())
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

[DarrenZhang01]

Solved in 3b09697