Adding multichip tests

tenstorrent · Feb 5, 2025 · a059245 · a059245
1 parent f29e6c1
commit a059245
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/tests/infra/mulitchip_workload.py b/tests/infra/mulitchip_workload.py
@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass
+import jax
+from typing import Sequence
+
+from .workload import Workload
+
+
+@dataclass
+class MultichipWorkload(Workload):
+    """
+    Convenience dataclass storing a callable and its positional and keyword arguments.
+    """
+    mesh: jax.sharding.Mesh = None
+    in_specs: Sequence[jax.sharding.PartitionSpec] = None
+