[FEATURE]: Passing label to .viz()

[francescotaioli]

First of all, great project!

Describe the solution you'd like
Currently, when multiple viz() calls are made in sequence, it becomes challenging to differentiate between them.
To enhance clarity and usability, I propose allowing an additional parameter in the viz() function, such as a label, which can be used to annotate the visualizations.

Additional context

my_tensor = torch.randint(0,2, (2, 196), dtype=torch.bool)
my_tensor.viz(label="Random IDXs tensor")

...
# Another viz called with another label associated with it

[epistoteles]

Hi @francescotaioli, and thank you for the issue!

I really like the idea - I'm unhappy with the limitations that come with the current labeling in general. For example, tensors have labels by default, while images do not.

I was considering changing this so that by default nothing has labels and you can change the default by setting tensorhue.set_printoptions(label=True).

Another question that you correctly raise is that currently you can't customize your label. For example, in tensors -inf and +inf have their own special color. It could be nice to include that information in a legend next to the tensor. You could also think about adding the text you mentioned not as label (below the tensor) but as title (above the tensor). There is also room right of the tensor.

Do you mind sharing your thoughts on this? What would you prefer?

[francescotaioli]

I was considering changing this so that by default nothing has labels and you can change the default by setting

This sounds like a great idea, as long as the label argument remains part of the function definition (thus you could turn on for that particular function call)

legend next to the tensor.

It would be fantastic to include a legend for every tensor displayed. For instance, even for a boolean tensor, a legend could clarify which color corresponds to True and which to False.

You could also think about adding the text you mentioned not as label (below the tensor) but as title (above the tensor)

I was thinking the same thing—I may have used the wrong terminology earlier. The goal is to have a title (passed as an argument) that visually explains the tensor being displayed. You might also allow users to choose a specific color for this title—perhaps using a library like colorama for that purpose.

[epistoteles]

I'm a little conflicted when it comes to a title above the tensor - this is essentially just a regular print statement that you could prepend to your .viz() call, like here:

If you want to add color or other formatting, you could use rich (which is what powers TensorHue anyway), e.g. with:

from rich import print
print("[bold red]My tensor title[/bold red]"); torch.rand(40,40).viz()

Because of this I'm not sure what additional benefit one could derive from passing a title as an argument to viz() - unless it is integrated with the legend and in the same line.

What I believe could make sense is a lambda function as an argument to legend, e.g.

t.viz(legend=lambda tensor: f"shape = {tensor.shape}, dtype = {tensor.dtype}")

if you also want the dtype to be printed in the same line. Is this something that would help you?

This would also allow you to integrate arbitrary information from your tensor, including explanation strings that you add yourself with

t = torch.rand(40,40)
t.explanation = "my explanation"
torch.viz(legend=lambda tensor: "{tensor.explanation}, shape={tensor.shape}")

Can you tell me a little more about your use case, and let me know if that would already solve it?

[francescotaioli]

this is essentially just a regular print statement that you could prepend to your .viz() call

yes, but I think that an argument to viz would come in handy. In your example,
torch.rand(40,40).viz("Random tensor")

This could be displayed above the image, or in the legend.