How to support a dynamic shape input for tf.where() Tensorflow (tf.where) function does not support dynamic input shapes.

For example, the following function calculates the svd of a matrix A and tries to identify the singular values greater than a certain threshold to invert them while filling other values with zeros.

def func (A, threshold)
   s,u,v = tf.svd(A)
   condition = tf.greater(s, threshold)
   z= tf.where(condition , tf.reciprocal(s), tf.zeros(s.shape))
   return z


will result in an error:


  Cannot convert a partially known TensorShape to a Tensor: (?,)


Any idea how to overcome this issue?
Just do tf.zeros_like(s) instead of using the shape explicitly.