xla

67cb9ca6 - Optimize autocast tests. (#5970)

Optimize autocast tests. (#5970) * Only initialize once for the test suite instead of each test. * remove comments * removed unused lines * fix linter * fix a tpu issue * fix minor issue