PR #26279: Changed to math_ops.cast from math_ops.to_float

Imported from GitHub PR #26279

Switched to recommended method

Copybara import of the project:

  - 4c520c4511002b502704fa5788b9851dcd91b167 Changed to math_ops.cast from math_ops.to_float by Amit Srivastava <amit.srivastava@huawei.com>
  - fb4222d8238e273c5e4d7b17bec75ee1d9038e99 Changed to shape from get_shape() by Amit Srivastava <amit.srivastava@huawei.com>
  - 401ac3ad9176e2fbc743316845a6e2ed8b68b2cf Merge fb4222d8238e273c5e4d7b17bec75ee1d9038e99 into e31e0... by Amit <30853054+amitsrivastava78@users.noreply.github.com>

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/tensorflow/pull/26279 from amitsrivastava78:training_tc_1 fb4222d8238e273c5e4d7b17bec75ee1d9038e99
PiperOrigin-RevId: 236688845