Accept native Torch state in ModelCore and CoreSolver.run

.github/workflows/ci.yaml

@@ -15,6 +15,26 @@ permissions:
   contents: 'read'
 
 jobs:
+  torch-cpu:
+    name: 'PyTorch CPU (3.12)'
+    runs-on: 'ubuntu-latest'
+    steps:
+      - name: 'Checkout repository'
+        uses: 'actions/checkout@v6'
+        with:
+          persist-credentials: false
+      - name: 'Setup uv'
+        uses: 'astral-sh/setup-uv@v7.2.0'
+        with:
+          enable-cache: true
+          cache-dependency-glob: 'pyproject.toml'
+          python-version: '3.12'
+      - name: 'Install test dependencies and CPU PyTorch'
+        run: |
+          uv sync --dev
+          uv pip install 'torch>=2.6' --index-url https://download.pytorch.org/whl/cpu
+      - name: 'Run PyTorch contracts'
+        run: 'uv run --no-sync pytest tests/op_engine/test_array_namespace.py -k torch'
   build-check:
     name: 'Build Check (${{ matrix.python-version }})'
     runs-on: 'ubuntu-latest'

CHANGELOG.md

@@ -5,6 +5,13 @@ versioning while the project remains pre-1.0.
 
 ## [Unreleased]
 
+### Fixed
+
+- `ModelCore.set_initial_state` resolves its configured NumPy-compatible dtype
+  in the initial state's array namespace, so native PyTorch tensors work with
+  `CoreSolver.run`. CPU Torch CI covers fixed and adaptive explicit runs,
+  float32/float64, and optional history (#208).
+
 ## [0.6.0] - 2026-10-09
 
 ### Added

docs/guides/backends.md

@@ -12,10 +12,13 @@ Namespace portability and automatic differentiation are related but distinct.
 The Array API does not define `grad`, `jit`, or traced control flow. A namespace
 such as JAX can nevertheless differentiate the ordinary `op_engine` methods
 because their fixed-step numerical operations remain in that namespace.
-The optional `torch` dependency group qualifies representative eager
-fixed-step kernels with PyTorch autograd. Compiler integration and complete
-method qualification remain provider/backend capabilities rather than
-consequences of namespace selection.
+CPU PyTorch tests qualify `ModelCore` initialization and complete fixed and
+adaptive explicit runs with float32/float64 and optional history, plus a
+representative fixed-step kernel with PyTorch autograd. Configure the dtype
+using NumPy-compatible values such as `ModelCoreOptions(dtype=np.float32)`;
+`ModelCore` resolves that dtype in the initial state's namespace. Compiler
+integration and complete method qualification remain provider/backend
+capabilities rather than consequences of namespace selection.
 
 ## Portable fixed-step methods
 

src/op_engine/model_core.py

@@ -87,7 +87,8 @@ class ModelCoreOptions:
             This is metadata only and enables non-uniform grids and
             FV/FDM operator builders.
         store_history: Whether to store the full time history.
-        dtype: Floating-point dtype for internal arrays.
+        dtype: NumPy-compatible dtype for internal arrays, resolved to the
+            initial state's namespace when that state is supplied.
     """
 
     other_axes: tuple[int, ...] = ()
@@ -425,7 +426,7 @@ def set_initial_state(self, initial_state: Array) -> None:
         xp = _namespace_of(initial_state)
         initial_state_arr = cast(
             "Array",
-            xp.asarray(initial_state, dtype=self.dtype),
+            xp.asarray(initial_state, dtype=getattr(xp, self.dtype.name)),
         )
         if initial_state_arr.shape != self.state_shape:
             raise ValueError(

tests/op_engine/test_array_namespace.py

@@ -152,6 +152,59 @@ def test_model_core_options_no_longer_store_backend_module() -> None:
     assert "xp" not in {field.name for field in fields(ModelCoreOptions)}
 
 
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.parametrize("store_history", [False, True])
+@pytest.mark.parametrize(
+    ("method", "adaptive"),
+    [("rk4", False), ("dopri5", False), ("dopri5", True), ("heun", True)],
+)
+def test_core_solver_run_preserves_torch_state_and_matches_numpy(
+    dtype: type[np.float32 | np.float64],
+    method: str,
+    *,
+    adaptive: bool,
+    store_history: bool,
+) -> None:
+    """The complete run resolves NumPy dtype options into native Torch dtypes."""
+    torch = pytest.importorskip("torch")
+    times = np.asarray([0.0, 0.2, 0.7, 1.0])
+    initial = np.arange(1.0, 7.0).reshape(2, 3)
+    options = ModelCoreOptions(dtype=dtype, store_history=store_history)
+    numpy_core = ModelCore(2, 3, times, options=options)
+    torch_core = ModelCore(2, 3, times, options=options)
+    numpy_core.set_initial_state(initial)
+    torch_core.set_initial_state(torch.tensor(initial, dtype=torch.float64))
+
+    def rhs(_time: Scalar, state: Array) -> Array:
+        return cast("Array", array_namespace(state).multiply(state, -0.3))
+
+    config = RunConfig(
+        method=method,
+        adaptive=adaptive,
+        adaptive_cfg=AdaptiveConfig(rtol=1e-5, atol=1e-7, dt_init=0.1),
+    )
+    CoreSolver(numpy_core).run(rhs, config=config)
+    CoreSolver(torch_core).run(rhs, config=config)
+
+    result = torch_core.get_current_state()
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype == getattr(torch, np.dtype(dtype).name)
+    assert result.device.type == "cpu"
+    tolerance = 10 * np.finfo(dtype).eps
+    np.testing.assert_allclose(
+        result.numpy(), numpy_core.get_current_state(), rtol=tolerance
+    )
+    np.testing.assert_allclose(result.numpy(), initial * np.exp(-0.3), rtol=3e-6)
+    if store_history:
+        assert isinstance(torch_core.state_array, torch.Tensor)
+        assert torch_core.state_array.dtype == result.dtype
+        np.testing.assert_allclose(
+            torch_core.state_array.numpy(), numpy_core.state_array, rtol=tolerance
+        )
+    else:
+        assert torch_core.state_array is None
+
+
 def test_model_core_preserves_jax_namespace_for_updates_and_history() -> None:
     """Immutable JAX state/history updates never fall back to NumPy."""
     jnp = pytest.importorskip("jax.numpy")