chore: import upstream snapshot with attribution

tests/unit/inference/v2/model_implementations/parameters/utils.py

@@ -0,0 +1,60 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.allocator import on_device
+from deepspeed.inference.v2.inference_parameter import InferenceParameter
+from deepspeed.inference.v2.model_implementations.parameter_base import ParameterBase, ParametrizedList
+
+
+class SimpleParam(ParameterBase):
+    """
+    Parameter with single dependency.
+    """
+
+    param: torch.Tensor
+
+    @on_device
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform(self.param)
+
+
+class SimpleParametrizedList(ParametrizedList):
+    """
+    Parameter list based on `num_dependencies` attribute.
+    """
+
+    count_attr: str = "num_dependencies"
+
+
+class ListParam(ParameterBase):
+    """
+    Parameter with list dependency.
+
+    NOTE: This uses the tuple workaround for the `ParametrizedList` class
+    as described in the docstring of `ParametrizedList`.
+    """
+
+    params: SimpleParametrizedList
+
+    @on_device
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform(torch.cat(tuple(self.params)))
+
+
+class DummyInferenceModel:
+
+    @property
+    def num_dependencies(self) -> int:
+        return 2
+
+    def transform(self, param: torch.Tensor) -> torch.Tensor:
+        return InferenceParameter.initialize(param)
+
+
+def validate_device(tensor: torch.Tensor):
+    assert tensor.device == torch.device(get_accelerator().current_device())