ak_choose.py

def choose(a, choices, mode="raise", highlevel=True, behavior=None):
    def getfunction(inputs, depth):
        if not all(x.purelist_depth == 1 for x in inputs):
            return
        
        arrays = [nplike.asarray(x) for x in inputs]

        a = arrays[0]
        choices = arrays[1:]
        out = np.choose(a, choices, mode=mode)
        
        return lambda: (ak.layout.NumpyArray(out),)
    
    arrays = [a, *choices]
    inputs = [ak.to_layout(x) for x in arrays]
    nplike = ak.nplike.of(arrays)
    behavior = ak._util.behaviorof(arrays, behavior=behavior)
    
    out, = ak._util.broadcast_and_apply(inputs, getfunction, behavior)
    return ak._util.maybe_wrap(out, behavior, highlevel)