Clarify that lax.top_k returns values and corresponding indices in descending order

[Sohl-Dickstein]

Share the fruits of my debugging journey, so that future people like me are less likely to assume index order is preserved.

NOTE: The statement that values are sorted in descending order is based on observed behavior in a colab. I don't know if this is guaranteed or not.

[Sohl-Dickstein]

Yes please!

On Tue, Apr 16, 2024, 10:15 Peter Hawkins @.***> wrote:

@.**** commented on this pull request.

This is explicitly not something we promise: there are reasonable implementations (e.g., a heap, say) that don't have this property.

Would it suffice to say that explicitly: that order is not guaranteed?

— Reply to this email directly, view it on GitHub https://github.com/google/jax/pull/20781#pullrequestreview-2004212091, or unsubscribe https://github.com/notifications/unsubscribe-auth/AADZW4B6M7K6RRDMM2HWEYLY5VMBXAVCNFSM6AAAAABGJYFWP6VHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDAMBUGIYTEMBZGE . You are receiving this because you authored the thread.Message ID: @.***>

[Sohl-Dickstein]

I just made that change to the PR.

[jakevdp]

Before we merge this, could you squash your changes into a single commit? See https://jax.readthedocs.io/en/latest/contributing.html#single-change-commits-and-pull-requests

Thanks!

[Sohl-Dickstein]

Done. PTAL.