github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[tensor] added __repr__ to spec (#1147)

pull/1151/head
Frank Lee 2022-06-21 15:38:05 +08:00 committed by GitHub
parent 70dd88e2ee
commit 0e4e62d30d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/tensor/spec.py
View File

@ -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):
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}'