Automatic merge from submit-queue (batch tested with PRs 67375, 67353). If you want to cherry-pick this change to another branch, please follow the instructions <a href="https://github.com/kubernetes/community/blob/master/contributors/devel/cherry-picks.md">here</a>.
Allow passing timeout on TPU API operations via context
The caller will provide a context with timeout attached.
**Release note**:
```release-note
None
```
/assign @vishh
59fdc02b13