lovely-tensors 0.1.16

❤️ Lovely Tensors

❤️ Lovely Tensors

Read full docs

More lovely stuff

Working with numbers

• Numpy: ❤️ Lovely NumPy
• JAX: 💘 Lovely JAX
• TinyGrad: 🫀 Lovely Grad

Community

• Discord

Install

pip install lovely-tensors

or

mamba install lovely-tensors

or

conda install -c conda-forge lovely-tensors

How to use

How often do you find yourself debugging PyTorch code? You dump a tensor to the cell output, and see this:

numbers

tensor([[[-0.3541, -0.3369, -0.4054,  ..., -0.5596, -0.4739,  2.2489],
         [-0.4054, -0.4226, -0.4911,  ..., -0.9192, -0.8507,  2.1633],
         [-0.4739, -0.4739, -0.5424,  ..., -1.0390, -1.0390,  2.1975],
         ...,
         [-0.9020, -0.8335, -0.9363,  ..., -1.4672, -1.2959,  2.2318],
         [-0.8507, -0.7822, -0.9363,  ..., -1.6042, -1.5014,  2.1804],
         [-0.8335, -0.8164, -0.9705,  ..., -1.6555, -1.5528,  2.1119]],

        [[-0.1975, -0.1975, -0.3025,  ..., -0.4776, -0.3725,  2.4111],
         [-0.2500, -0.2325, -0.3375,  ..., -0.7052, -0.6702,  2.3585],
         [-0.3025, -0.2850, -0.3901,  ..., -0.7402, -0.8102,  2.3761],
         ...,
         [-0.4251, -0.2325, -0.3725,  ..., -1.0903, -1.0203,  2.4286],
         [-0.3901, -0.2325, -0.4251,  ..., -1.2304, -1.2304,  2.4111],
         [-0.4076, -0.2850, -0.4776,  ..., -1.2829, -1.2829,  2.3410]],

        [[-0.6715, -0.9853, -0.8807,  ..., -0.9678, -0.6890,  2.3960],
         [-0.7238, -1.0724, -0.9678,  ..., -1.2467, -1.0201,  2.3263],
         [-0.8284, -1.1247, -1.0201,  ..., -1.2641, -1.1596,  2.3786],
         ...,
         [-1.2293, -1.4733, -1.3861,  ..., -1.5081, -1.2641,  2.5180],
         [-1.1944, -1.4559, -1.4210,  ..., -1.6476, -1.4733,  2.4308],
         [-1.2293, -1.5256, -1.5081,  ..., -1.6824, -1.5256,  2.3611]]])

Was it really useful for you, as a human, to see all these numbers?

What is the shape? The size?
What are the statistics?
Are any of the values nan or inf?
Is it an image of a man holding a tench?

import lovely_tensors as lt

lt.monkey_patch()

Summary

numbers # torch.Tensor

tensor[3, 196, 196] n=115248 (0.4Mb) x∈[-2.118, 2.640] μ=-0.388 σ=1.073

Better, huh?

numbers[1,:6,1] # Still shows values if there are not too many.

tensor[6] x∈[-0.443, -0.197] μ=-0.311 σ=0.091 [-0.197, -0.232, -0.285, -0.373, -0.443, -0.338]

spicy = numbers[0,:12,0].clone()

spicy[0] *= 10000
spicy[1] /= 10000
spicy[2] = float('inf')
spicy[3] = float('-inf')
spicy[4] = float('nan')

spicy = spicy.reshape((2,6))
spicy # Spicy stuff

tensor[2, 6] n=12 x∈[-3.541e+03, -4.054e-05] μ=-393.842 σ=1.180e+03 +Inf! -Inf! NaN!

torch.zeros(10, 10) # A zero tensor - make it obvious

tensor[10, 10] n=100 all_zeros

spicy.v # Verbose

tensor[2, 6] n=12 x∈[-3.541e+03, -4.054e-05] μ=-393.842 σ=1.180e+03 +Inf! -Inf! NaN!
tensor([[-3.5405e+03, -4.0543e-05,         inf,        -inf,         nan, -6.1093e-01],
        [-6.1093e-01, -5.9380e-01, -5.9380e-01, -5.4243e-01, -5.4243e-01, -5.4243e-01]])

spicy.p # The plain old way

tensor([[-3.5405e+03, -4.0543e-05,         inf,        -inf,         nan, -6.1093e-01],
        [-6.1093e-01, -5.9380e-01, -5.9380e-01, -5.4243e-01, -5.4243e-01, -5.4243e-01]])

Named dimensions

named_numbers = numbers.rename("C", "H","W")
named_numbers

tensor[C=3, H=196, W=196] n=115248 (0.4Mb) x∈[-2.118, 2.640] μ=-0.388 σ=1.073

Going .deeper

numbers.deeper

tensor[3, 196, 196] n=115248 (0.4Mb) x∈[-2.118, 2.640] μ=-0.388 σ=1.073
  tensor[196, 196] n=38416 x∈[-2.118, 2.249] μ=-0.324 σ=1.036
  tensor[196, 196] n=38416 x∈[-1.966, 2.429] μ=-0.274 σ=0.973
  tensor[196, 196] n=38416 x∈[-1.804, 2.640] μ=-0.567 σ=1.178

# You can go deeper if you need to
# And we can use `.deeper` with named dimensions.

named_numbers.deeper(2)

tensor[C=3, H=196, W=196] n=115248 (0.4Mb) x∈[-2.118, 2.640] μ=-0.388 σ=1.073
  tensor[H=196, W=196] n=38416 x∈[-2.118, 2.249] μ=-0.324 σ=1.036
    tensor[W=196] x∈[-1.912, 2.249] μ=-0.673 σ=0.522
    tensor[W=196] x∈[-1.861, 2.163] μ=-0.738 σ=0.418
    tensor[W=196] x∈[-1.758, 2.198] μ=-0.806 σ=0.397
    tensor[W=196] x∈[-1.656, 2.249] μ=-0.849 σ=0.369
    tensor[W=196] x∈[-1.673, 2.198] μ=-0.857 σ=0.357
    tensor[W=196] x∈[-1.656, 2.146] μ=-0.848 σ=0.372
    tensor[W=196] x∈[-1.433, 2.215] μ=-0.784 σ=0.397
    tensor[W=196] x∈[-1.279, 2.249] μ=-0.695 σ=0.486
    tensor[W=196] x∈[-1.364, 2.249] μ=-0.637 σ=0.539
    ...
  tensor[H=196, W=196] n=38416 x∈[-1.966, 2.429] μ=-0.274 σ=0.973
    tensor[W=196] x∈[-1.861, 2.411] μ=-0.529 σ=0.556
    tensor[W=196] x∈[-1.826, 2.359] μ=-0.562 σ=0.473
    tensor[W=196] x∈[-1.756, 2.376] μ=-0.622 σ=0.458
    tensor[W=196] x∈[-1.633, 2.429] μ=-0.664 σ=0.430
    tensor[W=196] x∈[-1.651, 2.376] μ=-0.669 σ=0.399
    tensor[W=196] x∈[-1.633, 2.376] μ=-0.701 σ=0.391
    tensor[W=196] x∈[-1.563, 2.429] μ=-0.670 σ=0.380
    tensor[W=196] x∈[-1.475, 2.429] μ=-0.616 σ=0.386
    tensor[W=196] x∈[-1.511, 2.429] μ=-0.593 σ=0.399
    ...
  tensor[H=196, W=196] n=38416 x∈[-1.804, 2.640] μ=-0.567 σ=1.178
    tensor[W=196] x∈[-1.717, 2.396] μ=-0.982 σ=0.350
    tensor[W=196] x∈[-1.752, 2.326] μ=-1.034 σ=0.314
    tensor[W=196] x∈[-1.648, 2.379] μ=-1.086 σ=0.314
    tensor[W=196] x∈[-1.630, 2.466] μ=-1.121 σ=0.305
    tensor[W=196] x∈[-1.717, 2.448] μ=-1.120 σ=0.302
    tensor[W=196] x∈[-1.717, 2.431] μ=-1.166 σ=0.314
    tensor[W=196] x∈[-1.560, 2.448] μ=-1.124 σ=0.326
    tensor[W=196] x∈[-1.421, 2.431] μ=-1.064 σ=0.383
    tensor[W=196] x∈[-1.526, 2.396] μ=-1.047 σ=0.417
    ...

Now in .rgb color

The important queston - is it our man?

numbers.rgb

Maaaaybe? Looks like someone normalized him.

in_stats = ( (0.485, 0.456, 0.406),     # mean 
             (0.229, 0.224, 0.225) )    # std

# numbers.rgb(in_stats, cl=True) # For channel-last input format
numbers.rgb(in_stats)

It’s indeed our hero, the Tenchman!

.plt the statistics

(numbers+3).plt

(numbers+3).plt(center="mean", max_s=1000)

(numbers+3).plt(center="range")

See the .chans

# .chans will map values betwen [-1,1] to colors.
# Make our values fit into that range to avoid clipping.
mean = torch.tensor(in_stats[0])[:,None,None]
std = torch.tensor(in_stats[1])[:,None,None]
numbers_01 = (numbers*std + mean)
numbers_01

tensor[3, 196, 196] n=115248 (0.4Mb) x∈[0., 1.000] μ=0.361 σ=0.248

numbers_01.chans

Let’s try with a Convolutional Neural Network

from torchvision.models import vgg11

features: torch.nn.Sequential = vgg11().features

# I saved the first 5 layers in "features.pt"
_ = features.load_state_dict(torch.load("../features.pt", weights_only=True), strict=False)

# Activatons of the second max pool layer of VGG11
acts = (features[:6](numbers[None])[0]/2) # /2 to reduce clipping
acts

tensor[128, 49, 49] n=307328 (1.2Mb) x∈[0., 12.508] μ=0.367 σ=0.634 grad DivBackward0

acts[:4].chans(cmap="coolwarm", scale=4)

Grouping

# Make 8 images with progressively higher brightness and stack them 2x2x2.
eight_images = (torch.stack([numbers]*8)
                    .add(torch.linspace(-3, 3, 8)[:,None,None,None])
                    .mul(torch.tensor(in_stats[1])[:,None,None])
                    .add(torch.tensor(in_stats[0])[:,None,None])
                    .clamp(0,1)
                    .view(2,2,2,3,196,196)
)
eight_images

tensor[2, 2, 2, 3, 196, 196] n=921984 (3.5Mb) x∈[0., 1.000] μ=0.411 σ=0.369

eight_images.rgb

# Weights of the second conv layer of VGG11
features[3].weight

Parameter containing:
Parameter[128, 64, 3, 3] n=73728 (0.3Mb) x∈[-0.783, 0.776] μ=-0.004 σ=0.065 grad

I want +/- 2σ to fall in the range [-1..1]

weights = features[3].weight.data
weights = weights / (2*2*weights.std()) # *2 because we want 2σ on both sides, so 4σ
# weights += weights.std() * 2
weights.plt

# Weights of the second conv layer (64ch -> 128ch) of VGG11,
# grouped per output channel.
weights.chans(frame_px=1, gutter_px=0)

It’s a bit hard to see. Scale up 10x, but onyl show the first 4 filters.

weights[:4].chans(frame_px=1, gutter_px=0, scale=10)

Options | Docs

from lovely_tensors import set_config, config, lovely, get_config

set_config(precision=1, sci_mode=True, color=False)
torch.tensor([1, 2, torch.nan])

tensor[3] μ=1.5e+00 σ=7.1e-01 NaN! [1.0e+00, 2.0e+00, nan]

set_config(precision=None, sci_mode=None, color=None) # None -> Reset to defaults

print(torch.tensor([1., 2]))
# Or with config context manager.
with config(sci_mode=True, precision=5):
    print(torch.tensor([1., 2]))

print(torch.tensor([1., 2]))

tensor[2] μ=1.500 σ=0.707 [1.000, 2.000]
tensor[2] μ=1.50000e+00 σ=7.07107e-01 [1.00000e+00, 2.00000e+00]
tensor[2] μ=1.500 σ=0.707 [1.000, 2.000]

Without .monkey_patch

lt.lovely(spicy)

tensor[2, 6] n=12 x∈[-3.541e+03, -4.054e-05] μ=-393.842 σ=1.180e+03 +Inf! -Inf! NaN!

lt.lovely(spicy, verbose=True)

tensor[2, 6] n=12 x∈[-3.541e+03, -4.054e-05] μ=-393.842 σ=1.180e+03 +Inf! -Inf! NaN!
tensor([[-3.5405e+03, -4.0543e-05,         inf,        -inf,         nan, -6.1093e-01],
        [-6.1093e-01, -5.9380e-01, -5.9380e-01, -5.4243e-01, -5.4243e-01, -5.4243e-01]])

lt.lovely(numbers, depth=1)

tensor[3, 196, 196] n=115248 (0.4Mb) x∈[-2.118, 2.640] μ=-0.388 σ=1.073
  tensor[196, 196] n=38416 x∈[-2.118, 2.249] μ=-0.324 σ=1.036
  tensor[196, 196] n=38416 x∈[-1.966, 2.429] μ=-0.274 σ=0.973
  tensor[196, 196] n=38416 x∈[-1.804, 2.640] μ=-0.567 σ=1.178

lt.rgb(numbers, in_stats)

lt.plot(numbers, center="mean")

lt.chans(numbers_01)

Matplotlib integration | Docs

numbers.rgb(in_stats).fig # matplotlib figure

(numbers*0.3+0.5).chans.fig # matplotlib figure

numbers.plt.fig.savefig('pretty.svg') # Save it

!file pretty.svg; rm pretty.svg

pretty.svg: SVG Scalable Vector Graphics image

Add content to existing Axes

fig = plt.figure(figsize=(8,3))
fig.set_constrained_layout(True)
gs = fig.add_gridspec(2,2)
ax1 = fig.add_subplot(gs[0, :])
ax2 = fig.add_subplot(gs[1, 0])
ax3 = fig.add_subplot(gs[1,1:])

ax2.set_axis_off()
ax3.set_axis_off()

numbers_01.plt(ax=ax1)
numbers_01.rgb(ax=ax2)
numbers_01.chans(ax=ax3);

torch.compile()

Just works.

def func(x):
    return x*2

if torch.__version__ >= "2.0":
    func = torch.compile(func)

func(torch.tensor([1,2,3]))

tensor[3] i64 x∈[2, 6] μ=4.000 σ=2.000 [2, 4, 6]