Fix device-aware dtype handling in identity, gradient functions (#2835)

Fix device-aware dtype handling in identity, gradient functions (#2835) #10096