Add step closures (#84300)
Ports over the step closure functionality from PyTorch/XLA to Lazy Tensor Core:
References:
https://github.com/pytorch/xla/blob/205ae574c0a24e092899ea8610c360f93f5d8142/torch_xla/core/xla_model.py#L852-L900
https://github.com/pytorch/xla/blob/205ae574c0a24e092899ea8610c360f93f5d8142/torch_xla/utils/closures.py#L7-L83
CC: @wconstab @JackCaoG @Krovatkin
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84300
Approved by: https://github.com/JackCaoG, https://github.com/wconstab
