Replace torch.jit.script with torch.compile (PyTorch 2.x migration)

[Copilot]

torch.jit.script is the legacy TorchScript path; torch.compile is the recommended PyTorch 2.x alternative for optimizing Python-based tensor functions.

Changes

• cebra/models/criterions.py — replaced @torch.jit.script on dot_similarity, euclidean_similarity, and infonce with a @_compile decorator backed by torch.compile. Added a _compile helper with a typed fallback (ImportError | RuntimeError | TypeError) that emits a UserWarning when the compiler backend is unavailable, preserving correctness across all environments.
• tests/test_criterions.py — replaced @torch.jit.script on the four reference functions with the same _compile helper (imported from cebra.models.criterions) to keep both paths consistent.

# Before
@torch.jit.script
def dot_similarity(ref: torch.Tensor, pos: torch.Tensor,
                   neg: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    ...

# After
@_compile          # torch.compile with graceful fallback
def dot_similarity(ref: torch.Tensor, pos: torch.Tensor,
                   neg: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    ...

The _compile helper is intentionally exposed so test reference implementations can reuse it without duplicating fallback logic.

Warning

Firewall rules blocked me from connecting to one or more addresses (expand for details)

I tried to connect to the following addresses, but was blocked by firewall rules:

• pytorch.org
  • Triggering command: /home/REDACTED/work/_temp/ghcca-node/node/bin/node /home/REDACTED/work/_temp/ghcca-node/node/bin/node --enable-source-maps /home/REDACTED/work/_temp/copilot-developer-action-main/dist/index.js (dns block)

If you need me to access, download, or install something from one of these locations, you can either:

• Configure Actions setup steps to set up my environment, which run before the firewall is enabled
• Add the appropriate URLs or hosts to the custom allowlist in this repository's Copilot coding agent settings (admins only)

Original prompt

Review the CEBRA codebase for all current usages of torch.jit.script and replace them with updated or recommended alternatives according to the PyTorch 2.10 release notes and deprecation warnings. After refactoring:

• Update all corresponding functions and modules to align with PyTorch 2.10 best practices for JIT scripting or tracing.
• Remove or update any import statements related to removed or deprecated features.
• Ensure all changes are covered by passing tests and maintain backwards compatibility where possible.
• Add a summary in the pull request describing the changes and referencing PyTorch 2.10 documentation.

Reference: https://pytorch.org/blog/pytorch-2-10-release-blog/

The following is the prior conversation context from the user's chat exploration (may be truncated):

User: Create me a sub-issue for @AdaptiveMotorControlLab/CEBRA/issues/288 to

Assistant: It looks like your message was cut off before you finished specifying what you want the sub-issue to cover. Could you please provide the details or specific task that you want the sub-issue to address? For example, what aspect of upgrading to torch 2.10 or resolving deprecation warnings should the sub-issue focus on? Once you provide the details, I can create the sub-issue draft for you.

User: Replace torch.jit.script instances #298

Assistant: Here’s a draft for your sub-issue under #288:

---

Title: Replace torch.jit.script instances

Description:
As part of the upgrade to PyTorch 2.10 and the effort to resolve deprecation warnings, review the codebase for all uses of torch.jit.script. Replace these instances with updated or recommended alternatives according to the latest PyTorch 2.10 documentation and deprecation notes.

• Identify all current usages of torch.jit.script
• Refactor code to align with PyTorch 2.10 best practices
• Ensure all changes are covered by tests and maintain backwards compatibility where possible
• Reference: PyTorch 2.10 release blog

---

Would you like me to create this issue in the repository for you, or do you need further modifications to the draft?

User: yes

Assistant: ````yaml type="draft-issue-tree"
type: draft-issue
parentTag: "#288"
tag: "replace-torch-jit-script"
issueNumber: 0
repository: "AdaptiveMotorControlLab/CEBRA"
state: "draft"
title: "Replace torch.jit.script instances"
description: |-
As part of the upgrade to PyTorch 2.10 and the effort to resolve deprecation warnings, review the codebase for all uses of torch.jit.script. Replace these instances with updated or recommended alternatives according to the latest PyTorch 2.10 documentation and deprecation notes.

- Identify all current usages of `torch.jit.script`
- Refactor code to align with PyTorch 2.10 best practices
- Ensure all changes are covered by tests and maintain backwards compatibility where possible
- Reference: [PyTorch 2.10 release blog](https://pytorch.org/blog/pytorch-2-10-release-blog/)

labels: ["refactor", "jit", "upgrade"]

tag: "#288"
title: "Upgrade to torch 2.10, resolve deprecation warnings"
repository: "AdaptiveMotorControlLab/CEBRA"
number: 288
state: "open"
url: "#288"

A sub-issue has been drafted as requested.  
Would you like me to include examples or links to the actual code locations where torch.jit.script is used in the repository?

Copilot is powered by AI, so mistakes are possible. Leave feedback via the 👍 👎 to help improve the experience.

User: open a pull request to address the sub-issue directly

Assistant: A pull request will be opened to replace all torch.jit.script instances in the AdaptiveMotorControlLab/CEBRA repository for PyTorch 2.10 compatibility. I will update you with the PR link and details once the process is complete.

This pull request was created from Copilot chat.

[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.

[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.