🐛[BUG]: Can't use Transolver model without transformer_engine

[bstaber]

Version

1.3.0

On which installation method(s) does this occur?

Pip

Describe the issue

It seems like the Transolver model was intended to work even when transformer_engine is not installed. For example, in transolver.py, the import is wrapped in a try/except:

# physicsnemo/models/transolver/transolver.py
try:
    import transformer_engine.pytorch as te

    TE_AVAILABLE = True
except ImportError:
    TE_AVAILABLE = False

However, in Physics_Attention.py, transformer_engine is imported without a try/except:

# physicsnemo/models/transolver/Physics_Attention.py
import torch
import torch.nn as nn
import transformer_engine.pytorch as te  # noqa: F401
from einops import rearrange
from torch.autograd.profiler import record_function
from torch.distributed.tensor.placement_types import Replicate

Is this intentional, or should this import also be optional to allow using the model without transformer_engine installed?

Thanks!

Minimum reproducible example

I use this pyproject.toml for uv:

[project]
name = "sciml"
version = "0.1.0"
description = "Super description."
readme = "README.md"
requires-python = ">=3.11,<3.14"
dependencies = [
    "einops>=0.8.1",
    "mkdocs>=1.6.1",
    "mkdocs-include-markdown-plugin>=7.2.0",
    "mkdocs-material>=9.7.0",
    "mkdocs-section-index>=0.3.10",
    "mkdocstrings-python>=1.19.0",
    "mlflow>=3.6.0",
    "neuraloperator>=2.0.0",
    "nvidia-physicsnemo>=1.3.0",
    "open3d>=0.19.0",
    "pyplaid>=0.1.10",
    "rich>=14.2.0",
    "torch>=2.8.0,<2.9.0",
    "torch-scatter",
]

[build-system]
requires = ["uv_build>=0.8.18,<0.9.0"]
build-backend = "uv_build"

[[tool.uv.index]]
name = "pyg"
url = "https://data.pyg.org/whl/torch-2.8.0+cu129.html"
ignore-error-codes = [403]
format = "flat"

[tool.uv.sources]
torch-scatter = { index = "pyg" }

And then simply try:

from physicsnemo.models.transolver import Transolver

Relevant log output

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/brian/softs/sciml/.venv/lib/python3.12/site-packages/physicsnemo/models/transolver/__init__.py", line 33, in <module>
    from .transolver import Transolver
  File "/home/brian/softs/sciml/.venv/lib/python3.12/site-packages/physicsnemo/models/transolver/transolver.py", line 53, in <module>
    from .Physics_Attention import (
  File "/home/brian/softs/sciml/.venv/lib/python3.12/site-packages/physicsnemo/models/transolver/Physics_Attention.py", line 37, in <module>
    import transformer_engine.pytorch as te  # noqa: F401
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'transformer_engine'
>>> from physicsnemo.models.transolver import Transolver
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/brian/softs/sciml/.venv/lib/python3.12/site-packages/physicsnemo/models/transolver/__init__.py", line 33, in <module>
    from .transolver import Transolver
  File "/home/brian/softs/sciml/.venv/lib/python3.12/site-packages/physicsnemo/models/transolver/transolver.py", line 53, in <module>
    from .Physics_Attention import (
  File "/home/brian/softs/sciml/.venv/lib/python3.12/site-packages/physicsnemo/models/transolver/Physics_Attention.py", line 37, in <module>
    import transformer_engine.pytorch as te  # noqa: F401
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'transformer_engine'

Environment details

[coreyjadams]

Hi @bstaber ,

Thanks for reporting this. We've got this fixed in an upcoming merge, which should make transolver work without transformer_engine installed.

Should not be long until it's in the core repo, thanks!

Corey

[bstaber]

Hi @coreyjadams, thanks for the quick reply. I apologize, I should have checked PR #1242 first!