Fix some linting issues and test warnings.#260

shaobohou wants to merge 2 commits intomaingoogle-deepmind/tf2jax:mainfrom

fix_lintergoogle-deepmind/tf2jax:fix_linterCopy head branch name to clipboard