Some unit tests fail pip test
Fixed by https://github.com/tensorflow/tensorflow/commit/edc647632a2ec64e1d9a6a6dc699bad77e930738 Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60479">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60479">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

git HEAD

### Custom Code

No

### OS Platform and Distribution

Ubuntu 20.04

### Mobile device

n/a

### Python version

3.10

### Bazel version

5.3.0

### GCC/Compiler version

10.2.1

### CUDA/cuDNN version

n/a

### GPU model and memory

n/a

### Current Behaviour?

The unit tests //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_2gpus, //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_cpu, //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_nccl_2gpus and //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_nccl_local_2gpus fail since commit https://github.com/tensorflow/tensorflow/commit/540a08ea4ddc93a632dd8583ddeae5be5ee75bff
See https://github.com/tensorflow/tensorflow/actions/runs/4865981223/jobs/8676977815#step:5:29731

### Standalone code to reproduce the issue

```shell
bazel test --build_tests_only --config=mkl_aarch64_threadpool --copt=-flax-vector-conversions --test_env=TF_ENABLE_ONEDNN_OPTS=1 --test_env=TF2_BEHAVIOR=1 --define=no_tensorflow_py_deps=true --test_lang_filters=py --flaky_test_attempts=3 --test_size_filters=small,medium --test_output=errors --verbose_failures=true --test_keep_going --notest_verbose_timeout_warnings --local_test_jobs=64 --test_tag_filters=-nopip,-no_pip,-oss_serial,-no_oss,-oss_excluded,-v1only,-benchmark-test,-no_aarch64,-no_oss_py38,-no_oss_py39,-no_oss_py310 -k -- //bazel_pip/tensorflow/... -//bazel_pip/tensorflow/compiler/tf2tensorrt/... -//bazel_pip/tensorflow/compiler/xrt/... -//bazel_pip/tensorflow/core/tpu/... -//bazel_pip/tensorflow/go/... -//bazel_pip/tensorflow/java/... -//bazel_pip/tensorflow/python/integration_testing/... -//bazel_pip/tensorflow/tools/toolchains/... -//bazel_pip/tensorflow/lite/... -//bazel_pip/tensorflow/python/kernel_tests/nn_ops:atrous_conv2d_test
```


### Relevant log output

```shell
FAIL: //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_2gpus (see /tmpfs/bazel_output/_bazel_ubuntu/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/testlogs/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test_2gpus/test_attempts/attempt_1.log)
FAIL: //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_2gpus (see /tmpfs/bazel_output/_bazel_ubuntu/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/testlogs/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test_2gpus/test_attempts/attempt_2.log)
FAIL: //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_2gpus (see /tmpfs/bazel_output/_bazel_ubuntu/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/testlogs/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test_2gpus/test.log)

FAILED: //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_2gpus (Summary)
      /tmpfs/bazel_output/_bazel_ubuntu/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/testlogs/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test_2gpus/test.log
      /tmpfs/bazel_output/_bazel_ubuntu/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/testlogs/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test_2gpus/test_attempts/attempt_1.log
      /tmpfs/bazel_output/_bazel_ubuntu/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/testlogs/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test_2gpus/test_attempts/attempt_2.log
INFO: From Testing //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_2gpus:
==================== Test output for //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_2gpus:
2023-05-02 23:40:14.254127: I tensorflow/core/util/port.cc:116] Experimental oneDNN custom operations are on. If you experience issues, please turn them off by setting the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
Traceback (most recent call last):
  File "/tmpfs/bazel_output/_bazel_ubuntu/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test_2gpus.runfiles/org_tensorflow/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test.py", line 27, in <module>
    from tensorflow.dtensor.python.tests import multi_client_test_util
ImportError: cannot import name 'multi_client_test_util' from 'tensorflow.dtensor.python.tests' (/workspace/pip_test/venv_clean/lib/python3.10/site-packages/tensorflow/dtensor/python/tests/__init__.py)
================================================================================
==================== Test output for //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_2gpus:
2023-05-02 23:40:19.436998: I tensorflow/core/util/port.cc:116] Experimental oneDNN custom operations are on. If you experience issues, please turn them off by setting the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
Traceback (most recent call last):
  File "/tmpfs/bazel_output/_bazel_ubuntu/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test_2gpus.runfiles/org_tensorflow/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test.py", line 27, in <module>
    from tensorflow.dtensor.python.tests import multi_client_test_util
ImportError: cannot import name 'multi_client_test_util' from 'tensorflow.dtensor.python.tests' (/workspace/pip_test/venv_clean/lib/python3.10/site-packages/tensorflow/dtensor/python/tests/__init__.py)
================================================================================
==================== Test output for //bazel_pip/tensorflow/dtensor/python/tests:multi_client_test_2gpus:
2023-05-02 23:40:23.477863: I tensorflow/core/util/port.cc:116] Experimental oneDNN custom operations are on. If you experience issues, please turn them off by setting the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
Traceback (most recent call last):
  File "/tmpfs/bazel_output/_bazel_ubuntu/eab0d61a99b6696edb3d2aff87b585e8/execroot/org_tensorflow/bazel-out/aarch64-opt/bin/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test_2gpus.runfiles/org_tensorflow/bazel_pip/tensorflow/dtensor/python/tests/multi_client_test.py", line 27, in <module>
    from tensorflow.dtensor.python.tests import multi_client_test_util
ImportError: cannot import name 'multi_client_test_util' from 'tensorflow.dtensor.python.tests' (/workspace/pip_test/venv_clean/lib/python3.10/site-packages/tensorflow/dtensor/python/tests/__init__.py)
================================================================================
```
</details>