![]() 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_= possible_devices_= Ĭolocation group had the following types and supported devices: R"""Elementwise square root: :math:`\sqrt was colocated with a group of nodes that required incompatible device '/job:localhost/replica:0/task:0/device:GPU:0'. Do I understand correctly, that lax.sqrt implementation is not available in jax-metal plugin? How come so? Isn't it one of the most common operators?Īs next steps, I decided to try to patch lax source code myself, and replaced Then while investigating further, I tried to runĪnd got same error message. Loc("-":393:12): error: unregistered operation 'mps.square_root' found in dialect ('mps') that does not allow unknown operations I was about to try out training flax NN using jax-metal, when I was presented with
0 Comments
Leave a Reply. |
AuthorWrite something about yourself. No need to be fancy, just an overview. ArchivesCategories |