def update_numpydoc(docstr, fun, op):
    """Transforms the numpy docstring to remove references of
       parameters that are supported by the numpy version but not the JAX version"""
    if not hasattr(op, '__code__'):
        return docstr
    if docstr[:4] == '    ':
        lines = docstr.split('\n')
        for idx, line in enumerate(lines):
            lines[idx] = line.replace('    ', '', 1)
        docstr = '\n'.join(lines)
    begin_idx = docstr.find('Parameters')
    begin_idx = docstr.find('--\n', begin_idx) + 2
    end_idx = docstr.find('Returns', begin_idx)
    parameters = docstr[begin_idx:end_idx]
    param_list = parameters.replace('\n    ', '@@').split('\n')
    for idx, p in enumerate(param_list):
        param = p[:p.find(' : ')].split(', ')[0]
        if param not in op.__code__.co_varnames:
            param_list[idx] = ''
    param_list = [param for param in param_list if param != '']
    parameters = '\n'.join(param_list).replace('@@', '\n    ')
    return docstr[:begin_idx + 1] + parameters + docstr[end_idx - 2:]
