Fix device-aware dtype handling in identity, gradient functions

Fix device-aware dtype handling in identity, gradient functions #10085