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?

Hi @Joseph_Ruskiewicz ,

In tfjs, tf.where is similar to NumPy’s np.where , but there’s a key difference. np.where can accept a single condition or multiple conditions, while tf.where requires an array of conditions, one for each element.

According to my understanding, your map variable would likely contain a tf.tensor1d with values like [false, true, false, ...] . You then pass this as the condition to tf.where . For elements where the condition is true, tf.where evaluates trend.mul(2 * Math.PI).cos() . Otherwise, it evaluates tf.exp(trend.mul(3)) .

I wasn’t able to run your original code, but I made some changes and ran it successfully. Here’s the modified version:

const trend = [3, 2, 4, 3, 1, 3, 4];
const condition = trend.map( s => s < 0.4,)
const t2 = tf.where(condition,
                            tf.cos(tf.mul(trend, 2 * Math.PI)), 
                            tf.exp(tf.mul(trend, 3)));
t2.print();

output:

Let me know if it helps.

Thank You!!