Tensorflow 2.0 , replace 0 values in a tensor with 1s

Question

So I'm trying to implement a centering and scaling with tensorflow where I need to replace values==0 in the tensor with 1.0.

I know to do this with numpy x_std[x_std == 0.0] = 1.0 but can't figure out the best way to do it on tensorflow 2.0.

The code is:

def _center_scale_xy(X, Y, scale=True):
    """ Center X, Y and scale if the scale parameter==True
    Returns
    -------
        X, Y, x_mean, y_mean, x_std, y_std
    """
    # center
    x_mean = tf.reduce_mean(X,axis=0)
    X -= x_mean
    y_mean = tf.reduce_mean(Y,axis=0)
    Y -= y_mean
    # scale
    if scale:
        x_std = tf.math.reduce_std(X,axis=0)
        #x_std[x_std == 0.0] = 1.0 #This one I need to implement with tensors
        X /= x_std
        y_std = tf.math.reduce_std(Y,axis=0)
        y_std[y_std == 0.0] = 1.0
        Y /= y_std
    else:
        x_std = np.ones(X.shape[1])
        y_std = np.ones(Y.shape[1])
    return X, Y, x_mean, y_mean, x_std, y_std

Answer 1

Use tf.where like this:

x_std = tf.where(tf.equal(x_std, 0), tf.ones_like(x_std), x_std)