Make the test of the infinte_fcn example file pass.

[DarrenZhang01]

1. Add the tests for infinite_fcn in Travis CI;
2. Change uint32 to int32 for abstract key initialization for type compatibility reasons;
3. Update dataclasses with the PyTree node registration being removed;
4. Update eval_on_shapes and the according usage based on the latest update on Google Trax - https://github.com/google/trax/pull/970;
5. Remove the printing statements that were used for debugging purposes.