Autocasting now only calls expensive nest.map_structure when Tensors need to be
autocast. In the common case where Tensors are passed with the correct dtype,
minimal work is performed.

PiperOrigin-RevId: 312800528
Change-Id: I25cc00c3309ea48b6fdc5ce6915701b960907008