jax

903b50eb - Some cleanup and reformatting in `xla.py`.

Some cleanup and reformatting in `xla.py`. - Make creation of a few dictionaries more readable. - Use f-strings where possible. - Remove unused imports and function parameters. - Don't format string before passing to `log` function.