def _searchsorted(sorted_sequence, values, side='left', out_type=np.int32, name=None):
    'Find indices for insertion for list to remain sorted.'
    if JAX_MODE:
        try:
            func = _searchsorted_vmap_sides[side]
        except KeyError:
            raise ValueError(("'%s' is an invalid value for keyword 'side'" % side))
        sorted_sequence_2d = np.reshape(sorted_sequence, ((- 1), sorted_sequence.shape[(- 1)]))
        values_2d = np.reshape(values, ((- 1), values.shape[(- 1)]))
        if (sorted_sequence_2d.shape[0] != values_2d.shape[0]):
            raise ValueError('Leading dim_size of both tensors must match.')
        return np.reshape(func(sorted_sequence_2d, values_2d).astype(out_type), values.shape)
    sorted_sequence = sorted_sequence[..., np.newaxis, :]
    values = values[..., :, np.newaxis]
    if (side == 'left'):
        is_in_right_location = (sorted_sequence < values)
    elif (side == 'right'):
        is_in_right_location = (sorted_sequence <= values)
    return np.sum(is_in_right_location, axis=(- 1)).astype(out_type)
Find indices for insertion for list to remain sorted.
def _searchsorted(sorted_sequence, values, side='left', out_type=np.int32, name=None):
    
    if JAX_MODE:
        try:
            func = _searchsorted_vmap_sides[side]
        except KeyError:
            raise ValueError(("'%s' is an invalid value for keyword 'side'" % side))
        sorted_sequence_2d = np.reshape(sorted_sequence, ((- 1), sorted_sequence.shape[(- 1)]))
        values_2d = np.reshape(values, ((- 1), values.shape[(- 1)]))
        if (sorted_sequence_2d.shape[0] != values_2d.shape[0]):
            raise ValueError('Leading dim_size of both tensors must match.')
        return np.reshape(func(sorted_sequence_2d, values_2d).astype(out_type), values.shape)
    sorted_sequence = sorted_sequence[..., np.newaxis, :]
    values = values[..., :, np.newaxis]
    if (side == 'left'):
        is_in_right_location = (sorted_sequence < values)
    elif (side == 'right'):
        is_in_right_location = (sorted_sequence <= values)
    return np.sum(is_in_right_location, axis=(- 1)).astype(out_type)