[TPU] Set the compute_dtype with XLAFSDP #9032

Job	Run time
pl-cpu (macOS-11, lightning, 3.8, 1.11)	16m 23s
pl-cpu (macOS-11, lightning, 3.9, 1.12)	19m 25s
pl-cpu (macOS-11, lightning, 3.10, 1.13)	17m 50s
pl-cpu (ubuntu-20.04, lightning, 3.8, 1.11)	10m 45s
pl-cpu (ubuntu-20.04, lightning, 3.9, 1.12)	15m 9s
pl-cpu (ubuntu-20.04, lightning, 3.10, 1.13)	10m 44s
pl-cpu (windows-2022, lightning, 3.8, 1.11)	15m 15s
pl-cpu (windows-2022, lightning, 3.9, 1.12)	17m 25s
pl-cpu (windows-2022, lightning, 3.10, 1.13)	15m 39s
pl-cpu (macOS-11, lightning, 3.10, 2.0)	20m 10s
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.0)	12m 47s
pl-cpu (windows-2022, lightning, 3.10, 2.0)	14m 28s
pl-cpu (macOS-12, pytorch, 3.11, 2.0)	22m 21s
pl-cpu (ubuntu-22.04, pytorch, 3.11, 2.0)	11m 48s
pl-cpu (windows-2022, pytorch, 3.11, 2.0)	15m 0s
pl-cpu (macOS-11, lightning, 3.8, 1.11, oldest)	18m 56s
pl-cpu (ubuntu-20.04, lightning, 3.8, 1.11, oldest)	10m 49s
pl-cpu (windows-2022, lightning, 3.8, 1.11, oldest)	17m 45s
pl-cpu (macOS-11, pytorch, 3.8, 1.13)	18m 13s
pl-cpu (ubuntu-20.04, pytorch, 3.8, 1.13)	11m 37s
pl-cpu (windows-2022, pytorch, 3.8, 1.13)	16m 5s
	5h 28m 34s

Provide feedback