Intel Mac with AMD GPU user needs your help! "InvalidArgumentError: Cannot assign a device for operation"

Intel Mac with AMD Radeon Pro 5600M 8 GB in macOS 12.1 user is trying to run a model with jupyter, with successful setup yet fail to run the model with GPU. Details as follows:

Successfully install tensorflow-macos, and tensorflow-metal using the following code in Python 3.8.10

SYSTEM_VERSION_COMPAT=0 python -m pip install tensorflow-macos
SYSTEM_VERSION_COMPAT=0 python -m pip install tensorflow-metal 

and try to run a model. When setting up,

print("Num CPUs Available: ", len(tf.config.list_physical_devices('CPU')))
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

Num CPUs Available: 1 Num GPUs Available: 1 Both CPU and GPU is detected, and things seem to go well. However, following error shows when I execute the model.fit.

InvalidArgumentError: Cannot assign a device for operation sequential_15/sequential_14/random_flip_7/stateful_uniform_full_int/RngReadAndSkip: Could not satisfy explicit device specification '' because the node {{colocation_node sequential_15/sequential_14/random_flip_7/stateful_uniform_full_int/RngReadAndSkip}} was colocated with a group of nodes that required incompatible device '/job:localhost/replica:0/task:0/device:GPU:0'. All available devices [/job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:GPU:0]. 
Colocation Debug Info:
Colocation group had the following types and supported devices: 
Root Member(assigned_device_name_index_=2 requested_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' assigned_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' resource_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' supported_device_types_=[CPU] possible_devices_=[]
RngReadAndSkip: CPU 
_Arg: GPU CPU 

Colocation members, user-requested devices, and framework assigned devices, if any:
  sequential_15_sequential_14_random_flip_7_stateful_uniform_full_int_rngreadandskip_resource (_Arg)  framework assigned device=/job:localhost/replica:0/task:0/device:GPU:0
  sequential_15/sequential_14/random_flip_7/stateful_uniform_full_int/RngReadAndSkip (RngReadAndSkip) 
  sequential_15/sequential_14/random_flip_7/stateful_uniform_full_int_1/RngReadAndSkip (RngReadAndSkip) 

	 [[{{node sequential_15/sequential_14/random_flip_7/stateful_uniform_full_int/RngReadAndSkip}}]] [Op:__inference_train_function_18668]

pip package list is as follows:

---------------------------- ---------
absl-py                      1.0.0
anyio                        3.5.0
appnope                      0.1.2
argon2-cffi                  21.3.0
argon2-cffi-bindings         21.2.0
asttokens                    2.0.5
astunparse                   1.6.3
attrs                        21.4.0
Babel                        2.9.1
backcall                     0.2.0
black                        21.12b0
bleach                       4.1.0
cachetools                   4.2.4
certifi                      2021.10.8
cffi                         1.15.0
charset-normalizer           2.0.10
click                        8.0.3
cycler                       0.11.0
debugpy                      1.5.1
decorator                    5.1.1
defusedxml                   0.7.1
entrypoints                  0.3
executing                    0.8.2
flatbuffers                  2.0
fonttools                    4.28.5
gast                         0.4.0
google-auth                  2.3.3
google-auth-oauthlib         0.4.6
google-pasta                 0.2.0
grpcio                       1.43.0
h5py                         3.6.0
idna                         3.3
importlib-metadata           4.10.0
importlib-resources          5.4.0
ipykernel                    6.7.0
ipython                      8.0.0
ipython-genutils             0.2.0
jedi                         0.18.1
Jinja2                       3.0.3
json5                        0.9.6
jsonschema                   4.4.0
jupyter-client               7.1.1
jupyter-core                 4.9.1
jupyter-server               1.13.3
jupyterlab                   3.2.8
jupyterlab-pygments          0.1.2
jupyterlab-server            2.10.3
keras                        2.7.0
Keras-Preprocessing          1.1.2
kiwisolver                   1.3.2
libclang                     12.0.0
Markdown                     3.3.6
MarkupSafe                   2.0.1
matplotlib                   3.5.1
matplotlib-inline            0.1.3
mistune                      0.8.4
mypy-extensions              0.4.3
nbclassic                    0.3.5
nbclient                     0.5.10
nbconvert                    6.4.0
nbformat                     5.1.3
nest-asyncio                 1.5.4
notebook                     6.4.7
numpy                        1.22.1
oauthlib                     3.1.1
opt-einsum                   3.3.0
packaging                    21.3
pandocfilters                1.5.0
parso                        0.8.3
pathspec                     0.9.0
pexpect                      4.8.0
pickleshare                  0.7.5
Pillow                       9.0.0
pip                          21.3.1
platformdirs                 2.4.1
prometheus-client            0.12.0
prompt-toolkit               3.0.24
protobuf                     3.19.3
ptyprocess                   0.7.0
pure-eval                    0.2.1
pyasn1                       0.4.8
pyasn1-modules               0.2.8
pycparser                    2.21
Pygments                     2.11.2
pyparsing                    3.0.6
pyrsistent                   0.18.1
python-dateutil              2.8.2
pytz                         2021.3
pyzmq                        22.3.0
requests                     2.27.1
requests-oauthlib            1.3.0
rsa                          4.8
Send2Trash                   1.8.0
setuptools                   56.0.0
six                          1.15.0
sniffio                      1.2.0
stack-data                   0.1.4
tensorboard                  2.7.0
tensorboard-data-server      0.6.1
tensorboard-plugin-wit       1.8.1
tensorflow-estimator         2.7.0
tensorflow-io-gcs-filesystem 0.23.1
tensorflow-macos             2.7.0
tensorflow-metal             0.3.0
termcolor                    1.1.0
terminado                    0.12.1
testpath                     0.5.0
tomli                        1.2.3
tornado                      6.1
traitlets                    5.1.1
typing_extensions            4.0.1
urllib3                      1.26.8
wcwidth                      0.2.5
webencodings                 0.5.1
websocket-client             1.2.3
Werkzeug                     2.0.2
wheel                        0.37.1
wrapt                        1.13.3
zipp                         3.7.0

Much appreciate if anyone can share your insight over this instance.

Hi @alex2214,

The issue is caused due to a missing op implementation on the GPU for RngReadAndSkip. This is being worked on and I will update here once it is resolved. In the mean while unfortunately parts of the script using this op need to be run on the CPU to avoid the colocation error.

RngReadAndSkip is registered on the GPU in tensorflow-metal==0.5.1.

Intel Mac with AMD GPU user needs your help! "InvalidArgumentError: Cannot assign a device for operation"
 
 
Q