Fix example_jacobi_mpi import crash with public wp.Device (GH-1577)

[aryanputta]

Description

Closes #1577. warp/examples/distributed/example_jacobi_mpi.py crashes at import, before any MPI or GPU work:

File ".../example_jacobi_mpi.py", line 38, in <module>
    def calc_default_device(mpi_comm: "MPI.Comm") -> wp.context.Device:
                                                     ^^^^^^^^^^
AttributeError: module 'warp' has no attribute 'context'. Did you mean: 'constant'?

The return annotation references wp.context.Device, but Warp no longer exposes a public context attribute (Device is re-exported at the top level from warp/__init__.py). Because the file does not use from __future__ import annotations, the annotation is evaluated at function-definition time and the missing attribute aborts the import.

Changes

• Annotate calc_default_device() with the public wp.Device type instead of the non-public wp.context.Device.
• Add a ### Fixed CHANGELOG entry under Unreleased.

This follows the project guideline to import from warp in public-facing code rather than warp._src (wp.Device is warp._src.context.Device re-exported).

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes. (No new test: this is a one-line annotation fix to an example; the existing example now imports.)
• The documentation is up to date with these changes.
• CHANGELOG.md is updated under the Unreleased section.

Validation summary

Verified against the installed Warp public API that the fix is correct and the original is broken:

• python -c "import warp as wp; print(wp.Device); print(hasattr(wp, 'context'))" prints
  <class 'warp._src.context.Device'> and False, confirming wp.Device resolves and wp.context does not exist (the source of the AttributeError).
• Grepped the example: wp.context appeared only on line 38; no other references remain after the change.

Full mpirun execution needs a CUDA-aware MPI build and multiple GPUs, which is out of scope for this import-time annotation fix.

Bug fix

Reproduce without this PR (any platform, no GPU needed):

import warp as wp
wp.context.Device  # AttributeError: module 'warp' has no attribute 'context'

Summary by CodeRabbit

• Bug Fixes
  • Fixed an import-time crash in the distributed Jacobi example by using a supported device type reference.
• Documentation
  • Updated the changelog to note the fix.

[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: 3356e336-b131-4f9d-b488-5b52f701c6f5

📥 Commits

Reviewing files that changed from the base of the PR and between 3b50df9 and e7f4771.

📒 Files selected for processing (2)

• CHANGELOG.md
• warp/examples/distributed/example_jacobi_mpi.py

---

📝 Walkthrough

Walkthrough

The Jacobi MPI example now annotates calc_default_device with wp.Device instead of wp.context.Device, and the changelog adds a fixed entry for the import crash.

Changes

Jacobi MPI import fix

Layer / File(s)	Summary
Device annotation update warp/examples/distributed/example_jacobi_mpi.py, CHANGELOG.md	calc_default_device now returns wp.Device, and the changelog records the import crash fix for the example.

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 fix: replacing the import-crashing annotation with public wp.Device.
Linked Issues check	✅ Passed	The code change fixes the import-time AttributeError by using the public wp.Device type, matching issue #1577's requirement.
Out of Scope Changes check	✅ Passed	The changelog update is directly related to documenting the same fix and no unrelated changes are present.
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 a one-line import-time crash in example_jacobi_mpi.py caused by a stale reference to the private wp.context.Device type, which is no longer a public attribute of the warp module. The return annotation is updated to the equivalent public type wp.Device.

• warp/examples/distributed/example_jacobi_mpi.py: Changes the return annotation of calc_default_device() from wp.context.Device to wp.Device; confirmed via grep that no other wp.context references remain in the file.
• CHANGELOG.md: Adds a correctly placed ### Fixed entry linking to example_jacobi_mpi.py fails at import with AttributeError: module 'warp' has no attribute 'context' #1577.

Confidence Score: 5/5

Safe to merge — the change is a one-line annotation correction that restores a broken import with no runtime behavior change.

The fix replaces a private module path (wp.context.Device) with the documented public alias (wp.Device) in a single return annotation. A grep across the repository confirms no other wp.context.Device references remain. The CHANGELOG entry is correctly positioned and linked to the issue. There are no logic changes, no new dependencies, and the annotation is not evaluated at runtime.

No files require special attention.

Important Files Changed

Filename	Overview
warp/examples/distributed/example_jacobi_mpi.py	Single-line fix: replaces the private wp.context.Device return annotation with the public wp.Device alias, resolving an AttributeError crash at import time.
CHANGELOG.md	Adds a ### Fixed entry under Unreleased describing the import crash fix with a link to the original issue; entry is correctly placed at the top of the Fixed section.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["import warp as wp"] --> B{"wp.context exists?"}
    B -- "Before fix (No)" --> C["AttributeError at import\n'warp has no attribute context'"]
    B -- "After fix (N/A)" --> D["wp.Device evaluated\nas annotation string"]
    A --> E["calc_default_device(\n  mpi_comm\n) -> wp.Device"]
    E --> F["wp.get_cuda_device_count()"]
    F --> G["Return wp.get_cuda_device(local_rank)"]

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"}}}%%
flowchart TD
    A["import warp as wp"] --> B{"wp.context exists?"}
    B -- "Before fix (No)" --> C["AttributeError at import\n'warp has no attribute context'"]
    B -- "After fix (N/A)" --> D["wp.Device evaluated\nas annotation string"]
    A --> E["calc_default_device(\n  mpi_comm\n) -> wp.Device"]
    E --> F["wp.get_cuda_device_count()"]
    F --> G["Return wp.get_cuda_device(local_rank)"]

Loading

_{Reviews (1): Last reviewed commit: "Fix example_jacobi_mpi import crash with..." | Re-trigger Greptile}