How can achieve np.where using TensorFlow.js

Hi all. Coming from python we have np.where that takes a conditional expression and two tensors and gives back a tensor that has the value corresponding either the a or the b tensor depending on the value of the conditional. For example, I am working with this line of python code

 other_trend =  np.where(season_time < 0.4, np.cos(trend * 2 * np.pi), np.exp(3 * trend))

and I would like to get that to work in js.

In TensorFlow.js the tf.where takes a Tensor or TensorLike object and not a conditional expression. I have worked around this by using arraySync() and mapping over this to create a Tensor with boolean values corresponding to what I want from the condition. Like this

let map = tf.tensor1d(trend.arraySync().map( s => s < 0.4,))
let otherTrend = tf.where(map,trend.mul(2 * Math.PI).cos(), tf.exp(trend.mul(3)))

Is this the solution for simulating np.where in TensorFlow.js?