Skip to content

GitLab

"...images/git@developer.sourcefind.cn:OpenDAS/colossalai.git" did not exist on "08f2920ea8f54aeabc65330ca9f7c97317737da6"
Unverified Commit 0e4e62d3 authored by Frank Lee's avatar Frank Lee Committed by GitHub
Browse files

[tensor] added __repr__ to spec (#1147)

parent 70dd88e2
Show whitespace changes
Inline Side-by-side
Showing with 6 additions and 0 deletions
colossalai/tensor/spec.py
View file @ 0e4e62d3
...@@ -18,6 +18,9 @@ class ParallelAction(object): ...@@ -18,6 +18,9 @@ class ParallelAction(object):
self.compute_pattern = compute_pattern self.compute_pattern = compute_pattern
self.gather_out = gather_out self.gather_out = gather_out
def __repr__(self):
return f'compute pattern: {self.compute_pattern}, gather out: {self.gather_out}'
class TensorSpec(object): class TensorSpec(object):
""" """
...@@ -72,3 +75,6 @@ class TensorSpec(object): ...@@ -72,3 +75,6 @@ class TensorSpec(object):
def has_compute_pattern(self, compute_pattern: ComputePattern): def has_compute_pattern(self, compute_pattern: ComputePattern):
return self.parallel_action.compute_pattern == compute_pattern return self.parallel_action.compute_pattern == compute_pattern
def __repr__(self):
return f'parallel action: {self.parallel_action}, dist_spec: {self.dist_spec}'
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment