Lazy codegen change for xla (#76717)
Codegen change to enable PyTorch/XLA to generate the first op in https://github.com/pytorch/xla/pull/3544.
@bdhirsh @wconstab
Pull Request resolved: https://github.com/pytorch/pytorch/pull/76717
Approved by: https://github.com/Krovatkin
