code: def _where_leading_dims(mask, val1, val2):
      """Same as `np.where`, but broadcasting to rightmost dimensions."""
      desired_shape = [len(np.array(mask))] + [1] * (
          len(np.array(val2).shape) - 1
      )
      return np.where(np.reshape(mask, desired_shape), val1, val2)
  
imports: import collections
  import numpy
  from tensorflow_probability.python.experimental.auto_batching import instructions
label_desc: Use image node_id 2 for calling a global function with example usage: _where_leading_dims(mask, val1, val2) and returns: np