NotFoundError could not find registered transfer manager for platform Host -- check target linkage [Op:__inference__jit_compiled_convolution_op_26169] TPU-VM
Hi @innat ,

I am able to run the code on Colab by following TPU initialization code as per documentation and it executed fine as per attached [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/7165c07c2e944e347ec7f4360045bc57/61976_tpu.ipynb).

It seems the issue might be specific to Kaggle environment ?
 @SuryanarayanaY Thanks for checking. I also check in Colab TPU. For your concern, please check this comment https://github.com/keras-team/tf-keras/issues/655#issuecomment-1733864636  @innat , You mean to say it has problem with TPU VM and on colab as it is TPU node it works fine? I would like to hear from concerned team. Thanks! > You mean to say it has problem with TPU VM and on colab as it is TPU node it works fine?

Yes.


cc. @@djherbis  @SuryanarayanaY 

> I would like to hear from concerned team.

Could you like to mention the appropriate person (from tf team)?
 Any update?
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

2.12

### Custom code

Yes

### OS platform and distribution

_No response_

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

I was trying to use `groups` parameter from `keras.layers.Conv2D` API on TPU-VM device. While running on GPU works but on TPU, it doesn't if `groups > 1`.

### Standalone code to reproduce the issue

Full [Code.](https://www.kaggle.com/code/ipythonx/github-issue-655-keras-tpu-vm/notebook)

```python
input_shape = (4, 28, 28, 9)
x = tf.random.normal(input_shape)
y = tf.keras.layers.Conv2D(
    27, 3, activation='relu', input_shape=input_shape[1:], groups=3
)(x)
print(y.shape)
```


### Relevant log output

```yaml
---------------------------------------------------------------------------
NotFoundError                             Traceback (most recent call last)
Cell In[133], line 3
      1 input_shape = (4, 28, 28, 9)
      2 x = tf.random.normal(input_shape)
----> 3 y = tf.keras.layers.Conv2D(
      4     27, 3, activation='relu', input_shape=input_shape[1:], groups=3
      5 )(x)
      7 print(y.shape)

File /usr/local/lib/python3.8/site-packages/keras/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File /usr/local/lib/python3.8/site-packages/tensorflow/python/eager/execute.py:52, in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     50 try:
     51   ctx.ensure_initialized()
---> 52   tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     53                                       inputs, attrs, num_outputs)
     54 except core._NotOkStatusException as e:
     55   if name is not None:

NotFoundError: Exception encountered when calling layer 'conv2d_335' (type Conv2D).

could not find registered transfer manager for platform Host -- check target linkage [Op:__inference__jit_compiled_convolution_op_26236]

Call arguments received by layer 'conv2d_335' (type Conv2D):
  • inputs=tf.Tensor(shape=(4, 28, 28, 9), dtype=float32)
```
