Skip to content

Commit

Permalink
Add repr method on tensor subclass
Browse files Browse the repository at this point in the history
ghstack-source-id: 1adc8a1bb74725d7c9b1709a5df70973d2b60497
Pull Request resolved: #397
  • Loading branch information
tugsbayasgalan committed Jun 18, 2024
1 parent f5b6ec9 commit dcb8f56
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions torchao/dtypes/aqt.py
Original file line number Diff line number Diff line change
Expand Up @@ -480,6 +480,9 @@ def _apply_fn_to_data(self, fn):
self.scale_and_zero = fn(self.scale_and_zero)
return self

def __repr__(self):
return f"TensorCoreTiledAQTLayout(packed_weight={self.packed_weight}, scale_and_zero={self.scale_and_zero})"

@classmethod
def __torch_dispatch__(cls, func, types, args, kwargs):
kwargs = {} if kwargs is None else kwargs
Expand All @@ -495,6 +498,8 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
"""
args[0].transposed = not args[0].transposed
return return_and_correct_aliasing(func, args, kwargs, args[0])

breakpoint()

raise NotImplementedError(
f"TensorCoreTiledAQTLayout dispatch: attempting to run {func}, this is not supported"
Expand Down

0 comments on commit dcb8f56

Please sign in to comment.