Shortcuts

Source code for torch._tensor_str

# mypy: allow-untyped-defs
import contextlib
import dataclasses
import math
import textwrap
from typing import Any, Dict, Optional

import torch
from torch import inf


@dataclasses.dataclass
class __PrinterOptions:
    precision: int = 4
    threshold: float = 1000
    edgeitems: int = 3
    linewidth: int = 80
    sci_mode: Optional[bool] = None


PRINT_OPTS = __PrinterOptions()


# We could use **kwargs, but this will give better docs
[docs]def set_printoptions( precision=None, threshold=None, edgeitems=None, linewidth=None, profile=None, sci_mode=None, ): r"""Set options for printing. Items shamelessly taken from NumPy Args: precision: Number of digits of precision for floating point output (default = 4). threshold: Total number of array elements which trigger summarization rather than full `repr` (default = 1000). edgeitems: Number of array items in summary at beginning and end of each dimension (default = 3). linewidth: The number of characters per line for the purpose of inserting line breaks (default = 80). Thresholded matrices will ignore this parameter. profile: Sane defaults for pretty printing. Can override with any of the above options. (any one of `default`, `short`, `full`) sci_mode: Enable (True) or disable (False) scientific notation. If None (default) is specified, the value is defined by `torch._tensor_str._Formatter`. This value is automatically chosen by the framework. Example:: >>> # Limit the precision of elements >>> torch.set_printoptions(precision=2) >>> torch.tensor([1.12345]) tensor([1.12]) >>> # Limit the number of elements shown >>> torch.set_printoptions(threshold=5) >>> torch.arange(10) tensor([0, 1, 2, ..., 7, 8, 9]) >>> # Restore defaults >>> torch.set_printoptions(profile='default') >>> torch.tensor([1.12345]) tensor([1.1235]) >>> torch.arange(10) tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) """ if profile is not None: if profile == "default": PRINT_OPTS.precision = 4 PRINT_OPTS.threshold = 1000 PRINT_OPTS.edgeitems = 3 PRINT_OPTS.linewidth = 80 elif profile == "short": PRINT_OPTS.precision = 2 PRINT_OPTS.threshold = 1000 PRINT_OPTS.edgeitems = 2 PRINT_OPTS.linewidth = 80 elif profile == "full": PRINT_OPTS.precision = 4 PRINT_OPTS.threshold = inf PRINT_OPTS.edgeitems = 3 PRINT_OPTS.linewidth = 80 if precision is not None: PRINT_OPTS.precision = precision if threshold is not None: PRINT_OPTS.threshold = threshold if edgeitems is not None: PRINT_OPTS.edgeitems = edgeitems if linewidth is not None: PRINT_OPTS.linewidth = linewidth PRINT_OPTS.sci_mode = sci_mode
def get_printoptions() -> Dict[str, Any]: r"""Gets the current options for printing, as a dictionary that can be passed as ``**kwargs`` to set_printoptions(). """ return dataclasses.asdict(PRINT_OPTS) @contextlib.contextmanager def printoptions(**kwargs): r"""Context manager that temporarily changes the print options. Accepted arguments are same as :func:`set_printoptions`.""" old_kwargs = get_printoptions() set_printoptions(**kwargs) try: yield finally: set_printoptions(**old_kwargs) def tensor_totype(t): dtype = torch.float if t.is_mps else torch.double return t.to(dtype=dtype) class _Formatter: def __init__(self, tensor): self.floating_dtype = tensor.dtype.is_floating_point self.int_mode = True self.sci_mode = False self.max_width = 1 with torch.no_grad(): tensor_view = tensor.reshape(-1) if not self.floating_dtype: for value in tensor_view: value_str = f"{value}" self.max_width = max(self.max_width, len(value_str)) else: nonzero_finite_vals = torch.masked_select( tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0) ) if nonzero_finite_vals.numel() == 0: # no valid number, do nothing return # Convert to double for easy calculation. HalfTensor overflows with 1e8, and there's no div() on CPU. nonzero_finite_abs = tensor_totype(nonzero_finite_vals.abs()) nonzero_finite_min = tensor_totype(nonzero_finite_abs.min()) nonzero_finite_max = tensor_totype(nonzero_finite_abs.max()) for value in nonzero_finite_vals: if value != torch.ceil(value): self.int_mode = False break if self.int_mode: # in int_mode for floats, all numbers are integers, and we append a decimal to nonfinites # to indicate that the tensor is of floating type. add 1 to the len to account for this. if ( nonzero_finite_max / nonzero_finite_min > 1000.0 or nonzero_finite_max > 1.0e8 ): self.sci_mode = True for value in nonzero_finite_vals: value_str = f"{{:.{PRINT_OPTS.precision}e}}".format(value) self.max_width = max(self.max_width, len(value_str)) else: for value in nonzero_finite_vals: value_str = f"{value:.0f}" self.max_width = max(self.max_width, len(value_str) + 1) else: # Check if scientific representation should be used. if ( nonzero_finite_max / nonzero_finite_min > 1000.0 or nonzero_finite_max > 1.0e8 or nonzero_finite_min < 1.0e-4 ): self.sci_mode = True for value in nonzero_finite_vals: value_str = f"{{:.{PRINT_OPTS.precision}e}}".format(value) self.max_width = max(self.max_width, len(value_str)) else: for value in nonzero_finite_vals: value_str = f"{{:.{PRINT_OPTS.precision}f}}".format(value) self.max_width = max(self.max_width, len(value_str)) if PRINT_OPTS.sci_mode is not None: self.sci_mode = PRINT_OPTS.sci_mode def width(self): return self.max_width def format(self, value): if self.floating_dtype: if self.sci_mode: ret = f"{{:{self.max_width}.{PRINT_OPTS.precision}e}}".format(value) elif self.int_mode: ret = f"{value:.0f}" if not (math.isinf(value) or math.isnan(value)): ret += "." else: ret = f"{{:.{PRINT_OPTS.precision}f}}".format(value) else: ret = f"{value}" return (self.max_width - len(ret)) * " " + ret def _scalar_str(self, formatter1, formatter2=None): if formatter2 is not None: real_str = _scalar_str(self.real, formatter1) imag_str = (_scalar_str(self.imag, formatter2) + "j").lstrip() # handles negative numbers, +0.0, -0.0 if imag_str[0] == "+" or imag_str[0] == "-": return real_str + imag_str else: return real_str + "+" + imag_str else: return formatter1.format(self.item()) def _vector_str(self, indent, summarize, formatter1, formatter2=None): # length includes spaces and comma between elements element_length = formatter1.width() + 2 if formatter2 is not None: # width for imag_formatter + an extra j for complex element_length += formatter2.width() + 1 elements_per_line = max( 1, int(math.floor((PRINT_OPTS.linewidth - indent) / (element_length))) ) def _val_formatter(val, formatter1=formatter1, formatter2=formatter2): if formatter2 is not None: real_str = formatter1.format(val.real) imag_str = (formatter2.format(val.imag) + "j").lstrip() # handles negative numbers, +0.0, -0.0 if imag_str[0] == "+" or imag_str[0] == "-": return real_str + imag_str else: return real_str + "+" + imag_str else: return formatter1.format(val) if summarize and not PRINT_OPTS.edgeitems: # Deal with edge case that negative zero is zero data = ["..."] elif summarize and self.size(0) > 2 * PRINT_OPTS.edgeitems: data = ( [_val_formatter(val) for val in self[: PRINT_OPTS.edgeitems].tolist()] + [" ..."] + [_val_formatter(val) for val in self[-PRINT_OPTS.edgeitems :].tolist()] ) else: data = [_val_formatter(val) for val in self.tolist()] data_lines = [ data[i : i + elements_per_line] for i in range(0, len(data), elements_per_line) ] lines = [", ".join(line) for line in data_lines] return "[" + ("," + "\n" + " " * (indent + 1)).join(lines) + "]" # formatter2 is only used for printing complex tensors. # For complex tensors, formatter1 and formatter2 are the formatters for tensor.real # and tensor.imag respesectively def _tensor_str_with_formatter(self, indent, summarize, formatter1, formatter2=None): dim = self.dim() if dim == 0: return _scalar_str(self, formatter1, formatter2) if dim == 1: return _vector_str(self, indent, summarize, formatter1, formatter2) if summarize and self.size(0) > 2 * PRINT_OPTS.edgeitems: slices = ( [ _tensor_str_with_formatter( self[i], indent + 1, summarize, formatter1, formatter2 ) for i in range(0, PRINT_OPTS.edgeitems) ] + ["..."] + [ _tensor_str_with_formatter( self[i], indent + 1, summarize, formatter1, formatter2 ) for i in range(len(self) - PRINT_OPTS.edgeitems, len(self)) ] ) else: slices = [ _tensor_str_with_formatter( self[i], indent + 1, summarize, formatter1, formatter2 ) for i in range(0, self.size(0)) ] tensor_str = ("," + "\n" * (dim - 1) + " " * (indent + 1)).join(slices) return "[" + tensor_str + "]" def _tensor_str(self, indent): if self.numel() == 0: return "[]" if self.has_names(): # There are two main codepaths (possibly more) that tensor printing goes through: # - tensor data can fit comfortably on screen # - tensor data needs to be summarized # Some of the codepaths don't fully support named tensors, so we send in # an unnamed tensor to the formatting code as a workaround. self = self.rename(None) summarize = self.numel() > PRINT_OPTS.threshold if self._is_zerotensor(): self = self.clone() # handle the negative bit if self.is_neg(): self = self.resolve_neg() if self.dtype in [ torch.float16, torch.bfloat16, torch.float8_e5m2, torch.float8_e5m2fnuz, torch.float8_e4m3fn, torch.float8_e4m3fnuz, ]: self = self.float() if self.dtype is torch.complex32: self = self.cfloat() if self.dtype.is_complex: # handle the conjugate bit self = self.resolve_conj() real_formatter = _Formatter( get_summarized_data(self.real) if summarize else self.real ) imag_formatter = _Formatter( get_summarized_data(self.imag) if summarize else self.imag ) return _tensor_str_with_formatter( self, indent, summarize, real_formatter, imag_formatter ) else: formatter = _Formatter(get_summarized_data(self) if summarize else self) return _tensor_str_with_formatter(self, indent, summarize, formatter) def _add_suffixes(tensor_str, suffixes, indent, force_newline): tensor_strs = [tensor_str] last_line_len = len(tensor_str) - tensor_str.rfind("\n") + 1 for suffix in suffixes: suffix_len = len(suffix) if force_newline or last_line_len + suffix_len + 2 > PRINT_OPTS.linewidth: tensor_strs.append(",\n" + " " * indent + suffix) last_line_len = indent + suffix_len force_newline = False else: tensor_strs.append(", " + suffix) last_line_len += suffix_len + 2 tensor_strs.append(")") return "".join(tensor_strs) def get_summarized_data(self): dim = self.dim() if dim == 0: return self if dim == 1: if self.size(0) > 2 * PRINT_OPTS.edgeitems: return torch.cat( (self[: PRINT_OPTS.edgeitems], self[-PRINT_OPTS.edgeitems :]) ) else: return self if not PRINT_OPTS.edgeitems: return self.new_empty([0] * self.dim()) elif self.size(0) > 2 * PRINT_OPTS.edgeitems: start = [self[i] for i in range(0, PRINT_OPTS.edgeitems)] end = [self[i] for i in range(len(self) - PRINT_OPTS.edgeitems, len(self))] return torch.stack([get_summarized_data(x) for x in (start + end)]) else: return torch.stack([get_summarized_data(x) for x in self]) def _str_intern(inp, *, tensor_contents=None): if torch._C._functorch.is_functorch_wrapped_tensor(inp): return _functorch_wrapper_str_intern(inp, tensor_contents=tensor_contents) is_plain_tensor = type(inp) is torch.Tensor or type(inp) is torch.nn.Parameter if inp.is_nested: prefix = "nested_tensor(" elif is_plain_tensor: prefix = "tensor(" else: prefix = f"{type(inp).__name__}(" indent = len(prefix) suffixes = [] custom_contents_provided = tensor_contents is not None if custom_contents_provided: tensor_str = tensor_contents # This is used to extract the primal value and thus disable the forward AD # within this function. # TODO(albanD) This needs to be updated when more than one level is supported self, tangent = torch.autograd.forward_ad.unpack_dual(inp) # Note [Print tensor device]: # A general logic here is we only print device when it doesn't match # the device specified in default tensor type. # Currently torch.set_default_tensor_type() only supports CPU/CUDA, thus # torch._C._get_default_device() only returns either cpu or cuda. # In other cases, we don't have a way to set them as default yet, # and we should always print out device for them. if ( self.device.type != torch._C._get_default_device() or ( self.device.type == "cuda" and torch.cuda.current_device() != self.device.index ) or (self.device.type == "mps") ): suffixes.append("device='" + str(self.device) + "'") # Tensor printing performs tensor operations like slice, indexing, etc to make it in a # representable format. These operations on ipu/xla/lazy/mtia tensor results in compilations. Hence, # to avoid compilations, copying the tensor to cpu before printing. if self.device.type in ["xla", "lazy", "ipu", "mtia"]: self = self.to("cpu") # TODO: add an API to map real -> complex dtypes _default_complex_dtype = ( torch.cdouble if torch.get_default_dtype() == torch.double else torch.cfloat ) has_default_dtype = self.dtype in ( torch.get_default_dtype(), _default_complex_dtype, torch.int64, torch.bool, ) if self.is_sparse: suffixes.append("size=" + str(tuple(self.shape))) from torch._subclasses.fake_tensor import FakeTensor is_meta = self.is_meta or isinstance(self, FakeTensor) if not is_meta: suffixes.append("nnz=" + str(self._nnz())) if not has_default_dtype: suffixes.append("dtype=" + str(self.dtype)) if not custom_contents_provided: indices_prefix = "indices=tensor(" indices = self._indices().detach() if is_meta: indices_str = "..." else: indices_str = _tensor_str(indices, indent + len(indices_prefix)) if is_meta or indices.numel() == 0: indices_str += ", size=" + str(tuple(indices.shape)) values_prefix = "values=tensor(" values = self._values().detach() if is_meta: values_str = "..." else: values_str = _tensor_str(values, indent + len(values_prefix)) if is_meta or values.numel() == 0: values_str += ", size=" + str(tuple(values.shape)) tensor_str = ( indices_prefix + indices_str + "),\n" + " " * indent + values_prefix + values_str + ")" ) elif self.layout in { torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc, }: from torch._subclasses.fake_tensor import FakeTensor suffixes.append("size=" + str(tuple(self.shape))) is_meta = self.is_meta or isinstance(self, FakeTensor) if not is_meta: suffixes.append("nnz=" + str(self._nnz())) if not has_default_dtype: suffixes.append("dtype=" + str(self.dtype)) if not custom_contents_provided: compressed_indices_method, plain_indices_method = { torch.sparse_csr: (torch.Tensor.crow_indices, torch.Tensor.col_indices), torch.sparse_csc: (torch.Tensor.ccol_indices, torch.Tensor.row_indices), torch.sparse_bsr: (torch.Tensor.crow_indices, torch.Tensor.col_indices), torch.sparse_bsc: (torch.Tensor.ccol_indices, torch.Tensor.row_indices), }[self.layout] if self.layout in {torch.sparse_csr, torch.sparse_bsr}: cdimname, pdimname = "row", "column" else: cdimname, pdimname = "column", "row" compressed_indices_prefix = f"c{cdimname[:3]}_indices=tensor(" compressed_indices = compressed_indices_method(self).detach() if is_meta: compressed_indices_str = "..." else: compressed_indices_str = _tensor_str( compressed_indices, indent + len(compressed_indices_prefix) ) if compressed_indices.numel() == 0 or is_meta: compressed_indices_str += ", size=" + str( tuple(compressed_indices.shape) ) plain_indices_prefix = f"{pdimname[:3]}_indices=tensor(" plain_indices = plain_indices_method(self).detach() if is_meta: plain_indices_str = "..." else: plain_indices_str = _tensor_str( plain_indices, indent + len(plain_indices_prefix) ) if plain_indices.numel() == 0 or is_meta: plain_indices_str += ", size=" + str(tuple(plain_indices.shape)) values_prefix = "values=tensor(" values = self.values().detach() if is_meta: values_str = "..." else: values_str = _tensor_str(values, indent + len(values_prefix)) if values.numel() == 0 or is_meta: values_str += ", size=" + str(tuple(values.shape)) tensor_str = ( compressed_indices_prefix + compressed_indices_str + "),\n" + " " * indent + plain_indices_prefix + plain_indices_str + "),\n" + " " * indent + values_prefix + values_str + ")" ) elif self.is_quantized: suffixes.append("size=" + str(tuple(self.shape))) if not has_default_dtype: suffixes.append("dtype=" + str(self.dtype)) suffixes.append("quantization_scheme=" + str(self.qscheme())) if ( self.qscheme() == torch.per_tensor_affine or self.qscheme() == torch.per_tensor_symmetric ): suffixes.append("scale=" + str(self.q_scale())) suffixes.append("zero_point=" + str(self.q_zero_point())) elif ( self.qscheme() == torch.per_channel_affine or self.qscheme() == torch.per_channel_symmetric or self.qscheme() == torch.per_channel_affine_float_qparams ): suffixes.append("scale=" + str(self.q_per_channel_scales())) suffixes.append("zero_point=" + str(self.q_per_channel_zero_points())) suffixes.append("axis=" + str(self.q_per_channel_axis())) if not custom_contents_provided: tensor_str = _tensor_str(self.dequantize(), indent) elif self.is_nested: if not custom_contents_provided: def indented_str(s, indent): return "\n".join(f" {line}" for line in s.split("\n")) strs = ",\n".join( indented_str(str(t), indent + 1) for t in torch.ops.aten.unbind.int(self, 0) ) tensor_str = f"[\n{strs}\n]" elif torch._is_functional_tensor(self): prefix = "_to_functional_tensor(" tensor_str = repr(torch._from_functional_tensor(self)) else: # Circular import problem, so we import it here from torch._subclasses.fake_tensor import FakeTensor if self.is_meta or isinstance(self, FakeTensor): suffixes.append("size=" + str(tuple(self.shape))) if self.dtype != torch.get_default_dtype(): suffixes.append("dtype=" + str(self.dtype)) # TODO: This implies that ellipses is valid syntax for allocating # a meta tensor or FakeTensor, which it could be, but it isn't right now if not custom_contents_provided: tensor_str = "..." else: if self.numel() == 0 and not self.is_sparse: # Explicitly print the shape if it is not (0,), to match NumPy behavior if self.dim() != 1: suffixes.append("size=" + str(tuple(self.shape))) # In an empty tensor, there are no elements to infer if the dtype # should be int64, so it must be shown explicitly. if self.dtype != torch.get_default_dtype(): suffixes.append("dtype=" + str(self.dtype)) if not custom_contents_provided: tensor_str = "[]" else: if not PRINT_OPTS.edgeitems: suffixes.append("size=" + str(tuple(self.shape))) if not has_default_dtype: suffixes.append("dtype=" + str(self.dtype)) if not custom_contents_provided: if self.layout != torch.strided: tensor_str = _tensor_str(self.to_dense(), indent) else: tensor_str = _tensor_str(self, indent) if self.layout != torch.strided: suffixes.append("layout=" + str(self.layout)) # Use inp here to get the original grad_fn and not the one generated by the forward grad # unpacking. grad_fn_name = None try: grad_fn = inp.grad_fn except RuntimeError: # Accessing the grad_fn calls rebasing logic which would cause an error # if that tensor is a view created in no-grad mode modified in-place in # no-grad mode. See: https://github.com/pytorch/pytorch/issues/99968 grad_fn_name = "Invalid" if grad_fn_name is None and grad_fn is not None: # type: ignore[possibly-undefined] grad_fn_name = type(grad_fn).__name__ if grad_fn_name == "CppFunction": grad_fn_name = grad_fn.name().rsplit("::", 1)[-1] if grad_fn_name is not None: suffixes.append(f"grad_fn=<{grad_fn_name}>") elif inp.requires_grad: suffixes.append("requires_grad=True") if self.has_names(): suffixes.append(f"names={self.names}") if tangent is not None: suffixes.append(f"tangent={tangent}") string_repr = _add_suffixes( prefix + tensor_str, suffixes, indent, force_newline=self.is_sparse # type: ignore[possibly-undefined] ) # Check if this instance is flagged as a parameter and change the repr accordingly. # Unfortunately, this function has to be aware of this detail. # NB: This is currently skipped for plain tensor parameters to maintain BC. In the future, # this should be done for those as well to produce a valid repr. if isinstance(self, torch.nn.Parameter) and not is_plain_tensor: string_repr = f"Parameter({string_repr})" return string_repr def _functorch_wrapper_str_intern(tensor, *, tensor_contents=None): level = torch._C._functorch.maybe_get_level(tensor) assert level != -1 if torch._C._functorch.is_functionaltensor(tensor): # Since we're unwrapping the FunctionalTensorWrapper, we need to make sure # that it's up to date first torch._sync(tensor) value = torch._C._functorch.get_unwrapped(tensor) value_repr = repr(value) indented_value_repr = textwrap.indent(value_repr, " " * 4) if torch._C._functorch.is_batchedtensor(tensor): bdim = torch._C._functorch.maybe_get_bdim(tensor) assert bdim != -1 return ( f"BatchedTensor(lvl={level}, bdim={bdim}, value=\n" f"{indented_value_repr}\n" f")" ) if torch._C._functorch.is_gradtrackingtensor(tensor): return ( f"GradTrackingTensor(lvl={level}, value=\n" f"{indented_value_repr}\n" f")" ) if torch._C._functorch.is_functionaltensor(tensor): return f"FunctionalTensor(lvl={level}, value=\\\n{value_repr})" raise ValueError("We don't know how to print this, please file us an issue") def _str(self, *, tensor_contents=None): with torch.no_grad(), torch.utils._python_dispatch._disable_current_modes(): guard = torch._C._DisableFuncTorch() return _str_intern(self, tensor_contents=tensor_contents)

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources