Jax support + updated testing code

• experimental support for Jax framework was added
• testing code was rewritten and updated to work
• tf2 always worked with einops, but tests had to be updated. So, tests are updated for tf2
• tf readme, minor additions, comments, etc.

Thanks to contributors