perf: Vectorize index clamping in Dataset.expand_index_in_trial

[SimpingOjou]

This PR addresses the TODOs in expand_index_in_trial regarding potential speed improvements. The previous implementation utilized a Python list comprehension to iterate over index and calculate boundaries, which created a bottleneck during execution.

Changes

• Removed the list comprehension used for boundary calculation and clamping.
• Implemented vectorized lookup for batch_trial_ids.
• Replaced the iterative clamp with a single operation over the entire tensor.

[cla-bot]

Thank you for your contribution. We require contributors to sign our Contributor License Agreement (CLA). We do not have a signed CLA on file for you. In order for us to review and merge your code, please sign our CLA here. After you signed, you can comment on this PR with @cla-bot check to trigger another check.

[SimpingOjou]

@cla-bot check

[cla-bot]

Thanks for tagging me. I looked for a signed form under your signature again, and updated the status on this PR. If the check was successful, no further action is needed. If the check was unsuccessful, please see the instructions in my first comment.

[stes]

Hi @SimpingOjou , thanks a lot for your contribution!

This looks good. Could you add a new test case, where you copy-paste the old implementation of the function to the tests as a backup, and then ensure that the results of the two implementations are the same?