Fix incorrect pairing in sample_plan_with_scipy

[zsteve]

What does this PR do?

Currently sample_plan_with_scipy returns incorrectly paired points due to an indexing bug. Should be

        pi_x0 = x0
        pi_x1 = x1[j]

not

       pi_x0 = x0[j]
       pi_x1 = x1

Before submitting

• Did you make sure title is self-explanatory and the description concisely explains the PR?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you list all the breaking changes introduced by this pull request?
• Did you test your PR locally with pytest command?
• Did you run pre-commit hooks with pre-commit run -a command?

Did you have fun?

Make sure you had fun coding 🙃

Summary by Sourcery

Bug Fixes:

• Correct point pairing in sample_plan_with_scipy by applying the assignment indices to the target points instead of the source points.

[sourcery-ai]

Reviewer's guide (collapsed on small PRs)

Reviewer's Guide

Fixes the point pairing logic in sample_plan_with_scipy by correcting which tensor is indexed by the assignment returned from SciPy’s linear_sum_assignment, so that x1 is permuted instead of x0.

Sequence diagram for corrected sample_plan_with_scipy pairing

sequenceDiagram
    participant Caller
    participant SamplePlan
    participant ScipyOptimize

    Caller->>SamplePlan: sample_plan_with_scipy(x0, x1)
    SamplePlan->>SamplePlan: compute cost matrix M from x0, x1
    SamplePlan->>SamplePlan: optionally normalize M
    SamplePlan->>ScipyOptimize: linear_sum_assignment(M)
    ScipyOptimize-->>SamplePlan: j
    SamplePlan->>SamplePlan: pi_x0 = x0
    SamplePlan->>SamplePlan: pi_x1 = x1[j]
    SamplePlan-->>Caller: (pi_x0, pi_x1)

Loading

File-Level Changes

Change	Details	Files
Correct the indexing in sample_plan_with_scipy so that the optimal assignment indices are applied to x1 instead of x0.	Keep x0 unchanged when forming the transport plan output Index x1 with the assignment indices j returned from linear_sum_assignment to produce correctly paired points Return the corrected (pi_x0, pi_x1) pair from sample_plan_with_scipy	torchcfm/optimal_transport.py

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey - I've reviewed your changes and they look great!

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[atong01]

Good catch! Thanks!