Skip to content

Use numpy C-API for C implementations of Sort/Argsort #1359

Open
@ricardoV94

Description

@ricardoV94

Description

These are available (for argsort maybe for a single axis) in the numpy public C-API:

https://numpy.org/devdocs/reference/c-api/array.html#c.PyArray_Sort

https://numpy.org/devdocs/reference/c-api/array.html#c.PyArray_ArgSort

Using it for our own ops should be simpler than the already simple SearchSorted implementation:

if self.mode == "full":
valid_conv = type(self)(mode="valid")
in1_bar = valid_conv(grad, in2[::-1])
in2_bar = valid_conv(grad, in1[::-1])

Sort may act inplace, so we can add an inplace version of the Sort op as well (and be sure to copy when not enabled)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

      Development

      No branches or pull requests

        Participants

        @ricardoV94

        Issue actions

          Use numpy C-API for C implementations of Sort/Argsort · Issue #1359 · pymc-devs/pytensor