Define Interpolate Tensorflow Function to maximize TPU performance

Question

How to perform interpolation vector with tensorflow that maximize TPU performance?

I have this code:

import tensorflow as tf

# Check if TPU is available
try:
  resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
  tf.config.experimental_connect_to_cluster(resolver)
  tf.tpu.experimental.initialize_tpu_system(resolver)
  strategy = tf.distribute.TPUStrategy(resolver)
  print("TPU initialized")
except:
  strategy = tf.distribute.get_strategy()
  print("No TPU detected, using default strategy for CPU/GPU")

# Define a custom TPU kernel function
@tf.function(jit_compile=True)
def interpolate_vector(uninterpolated_vector):
  interpolated_vector_condition = uninterpolated_vector != 0
  return interpolated_vector_condition
  return interpolated_vector # this should be return the vector, not condition

uninterpolated_vector = tf.constant([1,2,3,0,5,6,0,7,8,0,0,11], dtype=tf.float32)

# Use the custom TPU kernel within the TPU strategy scope
with strategy.scope():
  interpolated_vector = interpolate_vector(uninterpolated_vector)

print(uninterpolated_vector)
print(interpolated_vector)

Output:

tf.Tensor([ 1.  2.  3.  0.  5.  6.  0.  7.  8.  0.  0. 11.], shape=(12,), dtype=float32)
tf.Tensor([ True  True  True False  True  True False  True  True False False  True], shape=(12,), dtype=bool)

I can return condition vector what elements to interpolate, the condition is, if the element is zero, then it must be interpolated since there's no NULL, or None concept in dtype float32.

I expect it returns:

[ 1.  2.  3.  4.  5.  6.  6.5.  7.  8.  9.  10. 11.]

I also use @tf.function(jit_compile=True) to define and compile once of the function so that TPU accelerator can perform computing parallely faster withou evaluate function many times. Also I think I should avoid iteration programming (using for, while) since it will computing serially (sequencially) that depends on clock speed therefore it will be slower than parallel computation that depends on how much cores of TPU.

Since I avoid iteration programming, that's mean I need to define kernel function which is interpolate_vector. And I think this kernel will be elementwise kernel since it will return same shape as vector input.