Fix MapPeek crash when key is not a scaler
Also, please squash all commits, we don't need 3 commits for a check and a test. Thanks @cantonios. The PR has been updated with `with self.session():` removed. Let me know if this works (I see most of the other tests seems to have this) Thanks @mihaimaruseac  The PR has been updated with commits merged.  Please take a look.
 @mihaimaruseac  (In github there is an option in settings which allows the "squash merge" that can do the squash merge automatically when merge. This allows adjustments of try/revert tests in different sections while also keep the final commit history linear/clean.

Understand tensorflow has an internal merge system, though maybe something similar to "squash merge" can be considered as well?

_Below is a screenshot of the config in some other repos:_

-------

<img width="814" alt="screenshot" src="https://user-images.githubusercontent.com/6932348/218561562-4d8f417e-ae25-4f17-8f8b-37379863c978.png">

 That would be ideal, but unfortunately it requires significant changes in Copybara, so instead we're suggesting squashes on PRs if we can. This is useful for vuln fixes as it is easier to cherrypick a squashed commit than 3 commits on a fix.
This PR tries to address the issue raised in #58271 where MapPeek
will crash when key is not a scaler.

This PR fixes #58271.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>