Add lazy shape inference for take op (#82679)
Add lazy shape inference for take op
---
Companion PR on XLA's side: https://github.com/pytorch/xla/pull/3818
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82679
Approved by: https://github.com/JackCaoG
