fix(examples): use wp.Device in jacobi_mpi type hint

[drmikecrypto]

wp.context is not exposed on the public warp module; referencing wp.context.Device at import time raises AttributeError before MPI runs.

Fixes #1577

Description

Changes

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• CHANGELOG.md is updated for any user-facing changes under the Unreleased section.

Validation summary

Bug fix

import warp as wp
# Code that demonstrates the bug

New feature / enhancement

import warp as wp
# Code that demonstrates the new capability

Summary by CodeRabbit

• Chores
  • Updated the public type annotation in the distributed Jacobi example for consistency with the current device API.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Enterprise

Run ID: 21fd8d34-ea61-4a27-9dd2-4f695ced5516

📥 Commits

Reviewing files that changed from the base of the PR and between 6323e6d and b526cd0.

📒 Files selected for processing (1)

• warp/examples/distributed/example_jacobi_mpi.py

✅ Files skipped from review due to trivial changes (1)

• warp/examples/distributed/example_jacobi_mpi.py

---

📝 Walkthrough

Walkthrough

The MPI Jacobi example updates calc_default_device to return wp.Device instead of wp.context.Device, removing the import-time AttributeError.

Fix AttributeError on import

Layer / File(s)	Summary
Fix return type annotation warp/examples/distributed/example_jacobi_mpi.py	calc_default_device return annotation changed from wp.context.Device to wp.Device.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~2 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change to the Jacobi MPI example type hint.
Linked Issues check	✅ Passed	The change matches issue #1577 by replacing the import-time wp.context.Device annotation with wp.Device.
Out of Scope Changes check	✅ Passed	The PR only updates the public type hint needed to fix the import error and adds no unrelated changes.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[greptile-apps]

Greptile Summary

Fixes an AttributeError in the distributed Jacobi MPI example caused by referencing wp.context.Device as a return type hint — wp.context is not a public attribute of the warp module. The fix replaces it with wp.Device, which is the correct public export from warp/__init__.py.

• The return type annotation for calc_default_device is changed from wp.context.Device to wp.Device, matching the publicly exported symbol at warp/__init__.py:259.

Confidence Score: 5/5

Safe to merge — the change is a single-line type annotation fix that corrects an AttributeError without touching any runtime logic.

The only change is replacing wp.context.Device with wp.Device in a return type hint. wp.Device is confirmed to be exported from warp/__init__.py, so this is the correct public symbol. No runtime behaviour is altered, no tests are broken, and the fix directly addresses the reported import-time error.

No files require special attention.

Important Files Changed

Filename	Overview
warp/examples/distributed/example_jacobi_mpi.py	One-line fix: return type annotation updated from wp.context.Device to wp.Device to use the correct public API symbol

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Script as example_jacobi_mpi.py
    participant wp as warp (__init__.py)
    participant ctx as warp._src.context

    Script->>wp: import warp as wp
    wp->>ctx: from warp._src.context import Device as Device
    Note over Script,wp: wp.Device is now accessible (public export)
    Script->>wp: "calc_default_device() -> wp.Device (type hint)"
    Note over Script,ctx: wp.context is NOT exposed — using it would raise AttributeError

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Script as example_jacobi_mpi.py
    participant wp as warp (__init__.py)
    participant ctx as warp._src.context

    Script->>wp: import warp as wp
    wp->>ctx: from warp._src.context import Device as Device
    Note over Script,wp: wp.Device is now accessible (public export)
    Script->>wp: "calc_default_device() -> wp.Device (type hint)"
    Note over Script,ctx: wp.context is NOT exposed — using it would raise AttributeError

Loading

_{Reviews (2): Last reviewed commit: "fix(examples): use wp.Device in jacobi_m..." | Re-trigger Greptile}