Skip to content

Commit

Permalink
Some missing fixes left from previous PRs
Browse files Browse the repository at this point in the history
Summary:
* added repr for TensorCoreTiledAQTLayoutTensor: pytorch#397
* removed the str -> apply_tensor_class map that was used to implement string APIs, which is removed in pytorch#400

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:
  • Loading branch information
jerryzh168 committed Jun 25, 2024
1 parent 505edc1 commit 2dcc6ca
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 7 deletions.
4 changes: 4 additions & 0 deletions torchao/dtypes/affine_quantized_tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -480,6 +480,10 @@ def _apply_fn_to_data(self, fn):
self.scale_and_zero = fn(self.scale_and_zero)
return self

def __repr__(self):
int_data, scale, zero_point = self.get_plain()
return f"TensorCoreTiledAQTLayout(int_data={int_data}, scale={scale}, zero_point={zero_point})"

@classmethod
def __torch_dispatch__(cls, func, types, args, kwargs):
kwargs = {} if kwargs is None else kwargs
Expand Down
7 changes: 0 additions & 7 deletions torchao/quantization/quant_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -306,13 +306,6 @@ def filter_fn(module: nn.Module, fqn: str) -> bool:
m = quantize(m, apply_weight_quant, filter_fn)
"""
if isinstance(apply_tensor_subclass, str):
if apply_tensor_subclass not in _APPLY_TS_TABLE:
raise ValueError(f"{apply_tensor_subclass} not supported: {_APPLY_TS_TABLE.keys()}")
apply_tensor_subclass = _APPLY_TS_TABLE[apply_tensor_subclass]

assert not isinstance(apply_tensor_subclass, str)

_replace_with_custom_fn_if_matches_filter(
model,
_get_linear_subclass_inserter(apply_tensor_subclass),
Expand Down

0 comments on commit 2dcc6ca

Please sign in to comment.