fix(module): keep inputs on their device in a passthrough module without parameters

ChangeLog.md

@@ -7,6 +7,8 @@ SPDX-License-Identifier: Apache-2.0
 
 ## Unreleased
 
+- fix: a passthrough module without parameters keeps its inputs on their device instead of moving them to CPU.
+
 - feat (breaking): `inspect` measures module depth in executed modules.
   Containers such as `nn.ModuleList` and modules that never run no longer add a level.
   `min_depth` is renamed to `depth`.

tests/unit/torch/module/test_passthrough_module.py

@@ -6,9 +6,12 @@
 """Test for passhrough module."""
 
 import pytest
+import torch
+import torch.nn as nn
 
 from tests.toy_models import ToyTorchModel
 from torch_tweak.torch.module.passthrough_module import PassthroughModule
+from torch_tweak.torch.module.wrapper_module import Module
 from torch_tweak.torch.utils.xpu import is_available as is_xpu_available
 
 devices = ["cpu", "xpu"] if is_xpu_available() else ["cpu"]
@@ -21,3 +24,12 @@ def test_call(device):
     module = PassthroughModule(model, device=device)
 
     module(x)
+
+
+def test_module_without_parameters_keeps_inputs_on_their_device():
+    module = Module(nn.Flatten(), "flatten")
+    module.enable_passthrough()
+
+    output = module(torch.empty(2, 3, 4, device="meta"))
+
+    assert output.device.type == "meta"

torch_tweak/torch/module/passthrough_module.py

@@ -6,45 +6,39 @@
 
 from typing import Any
 
-from torch_tweak.torch.module.sample_metadata import SampleMetadata
+from torch_tweak.torch.module.locator import Locator
 
 
 class PassthroughModule:
     """Module that passes through the original module.
 
-    Before passing the data this module moves the data to the same device as the module.
-
+    When a device is given, inputs are moved to it before the call.
+    When no device is given, inputs stay where they are.
     """
 
     def __init__(
         self,
         module,
-        device: str,
+        device: str | None,
     ) -> None:
         """Initializes module.
 
         Args:
                 module: module to be tuned.
-                device: Device on which tuned module has to be executed.
+                device: Device on which tuned module has to be executed. None leaves the module and its
+                    inputs where they are, for modules without parameters.
         """
         super().__init__()
         self._forward_call = module.__call__
         self._device = device
-        module.to(self._device)
+        if self._device is not None:
+            module.to(self._device)
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         """Native inference on wrapped module."""
-        args, kwargs = self._prepare_inputs(*args, **kwargs)
+        if self._device is not None:
+            for locator, tensor in Locator.find_leaves(args, only_tensors=True):
+                args = locator.set_value(args, tensor.to(self._device))
+            for locator, tensor in Locator.find_leaves(kwargs, only_tensors=True):
+                kwargs = locator.set_value(kwargs, tensor.to(self._device))
         return self._forward_call(*args, **kwargs)
-
-    def _prepare_inputs(self, *args, **kwargs):
-        """Prepare inputs for inplace inference and place them on the same device if are not."""
-        sample_metadata = SampleMetadata.from_inputs(args, kwargs)
-
-        for locator, tensor_spec in sample_metadata.tensor_data:
-            if tensor_spec.name.startswith("args"):
-                args = locator.set_value(args, locator.get_value(args).to(self._device))
-            else:
-                kwargs = locator.set_value(kwargs, locator.get_value(kwargs).to(self._device))
-
-        return args, kwargs

torch_tweak/torch/module/wrapper_module.py

@@ -98,7 +98,7 @@ def __init__(
         module.forward = self._self_proxy_forward
         self._setup_strategies(strategy, strategies)
 
-        self._self_device = get_module_device(module) or torch.device("cpu")
+        self._self_device = get_module_device(module)
         self._self_state = ModuleState.INIT
         self._self_wrapper = None
         self._self_prev_recording = None
@@ -283,7 +283,7 @@ def tune(
         elif self._self_state == ModuleState.INIT or self._self_state == ModuleState.PASSTHROUGH:
             raise ValueError(f"Module: '{self._self_name}' has not recorded any samples. Cannot tune it.")
 
-        device = device or self._self_device
+        device = device or self._self_device or torch.device("cpu")
 
         recording = cast(RecordingModule, self._self_wrapper)
         backends: OrderedDict[SampleMetadata, Backend] = OrderedDict()