Skip to content

Commit 09e3869

Browse files
committed
Some missing fixes left from previous PRs
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:
1 parent 211b6bc commit 09e3869

File tree

2 files changed

+4
-7
lines changed

2 files changed

+4
-7
lines changed

torchao/dtypes/affine_quantized_tensor.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -480,6 +480,10 @@ def _apply_fn_to_data(self, fn):
480480
self.scale_and_zero = fn(self.scale_and_zero)
481481
return self
482482

483+
def __repr__(self):
484+
int_data, scale, zero_point = self.get_plain()
485+
return f"TensorCoreTiledAQTLayout(int_data={int_data}, scale={scale}, zero_point={zero_point})"
486+
483487
@classmethod
484488
def __torch_dispatch__(cls, func, types, args, kwargs):
485489
kwargs = {} if kwargs is None else kwargs

torchao/quantization/quant_api.py

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -310,13 +310,6 @@ def filter_fn(module: nn.Module, fqn: str) -> bool:
310310
"""
311311
if set_inductor_config:
312312
torchao.quantization.utils.recommended_inductor_config_setter()
313-
if isinstance(apply_tensor_subclass, str):
314-
if apply_tensor_subclass not in _APPLY_TS_TABLE:
315-
raise ValueError(f"{apply_tensor_subclass} not supported: {_APPLY_TS_TABLE.keys()}")
316-
apply_tensor_subclass = _APPLY_TS_TABLE[apply_tensor_subclass]
317-
318-
assert not isinstance(apply_tensor_subclass, str)
319-
320313
_replace_with_custom_fn_if_matches_filter(
321314
model,
322315
_get_linear_subclass_inserter(apply_tensor_subclass),

0 commit comments

Comments
 (0)