Here’s the Chainer implementation, if anyone’s curious. This could even be directly wrapped in a PyTorch function with a couple casting tricks (PyTorch tensors and cupy arrays have the same memory layout).

  
      github.com
  
  
    chainer/chainer/blob/master/chainer/functions/array/where.py 238
import numpy

from chainer.backends import cuda
from chainer import function_node
from chainer.utils import type_check


class Where(function_node.FunctionNode):

    """Choose elements depending on condition."""

    def check_type_forward(self, in_types):
        type_check.expect(in_types.size() == 3)
        c_type, x_type, y_type = in_types

        type_check.expect(
            c_type.dtype == numpy.bool_,
            x_type.dtype == y_type.dtype,
            x_type.shape == c_type.shape,
            y_type.shape == c_type.shape,


  This file has been truncated. show original

  
  
    
    
  
  


https://github.com/pfnet/chainer/blob/master/cupy/sorting/search.py#L92 36