個包装だんごとしてみる torch.Tensor シリーズの目次はこちら

2026-04-15 コードを追加しました。

この記事での約束

0 階のテンソル x.shape == () を、箱に入っていないだんごとします (スカラー)。
1 階のテンソル x.shape == (a,) を、a 個入りだんごの箱とします。
2 階のテンソル x.shape == (b, a) を、b 袋の a 個入りだんごの箱とします。
3 階のテンソル x.shape == (c, b, a) を、c 袋の b 袋の a 個入りだんごの箱とします。

なので、2 階以上のテンソルでは、だんごは袋に包装されています。
r 階のテンソルにおいて、だんごは r - 1 重に袋に包装されています。
袋のことを外側から袋 0 、袋 1 、...、袋 r - 2 と呼ぶことにします。

import torch

x = torch.tensor(1.)
assert x.shape == ()
print('箱に入っていないだんご\n', x)

x = torch.ones(4)
assert x.shape == (4,)
print('箱の中の 4 個入りだんご\n', x)

x = torch.ones(3, 4)
assert x.shape == (3, 4)
print('箱の中の 3 袋の 4 個入りだんご\n', x)

x = torch.ones(2, 3, 4)
assert x.shape == (2, 3, 4)
print('箱の中の 2 袋の 3 袋の 4 個入りだんご\n', x)

箱に入っていないだんご
 tensor(1.)

箱の中の 4 個入りだんご
 tensor([1., 1., 1., 1.])

箱の中の 3 袋の 4 個入りだんご
 tensor([[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.]])

箱の中の 2 袋の 3 袋の 4 個入りだんご
 tensor([[[1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.]],
         [[1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.]]])

x.mean

x.mean(dim=k, keepdim=True) は箱の中の袋 k たちを「平均的な袋 k」にまとめます。k = d - 1 のときは最も内側のだんごたちを「平均的なだんご」にまとめます。この結果、x.shape の k 階目が 1 に潰れます。

袋 k たちが「平均的袋 k」にまとめられるなら、もはや袋 k は 1 袋しかなく、袋 k は 1 袋目と 2 袋目とを分ける役割を果たしていません。そのため、こんな役割のない袋は残しておきたくないと思うなら keepdim=False にすれば袋をやぶり捨てることができます。k = d - 1 のときは「平均的だんご」が個包装になるのでその個包装が破り捨てられることになります。この場合は x.shape の k 階目が消えます。

実は後者の keepdim=False がデフォルト動作なのですが、計算して平均値でもって標準化するときなどは keepdim=True にしておくと元の箱とのブロードキャスト可能性 (後述) が保持されて便利です。

import torch

# 箱の中の 2 袋の 3 袋の 4 個入りだんご
x = torch.ones(2, 3, 4)
assert x.shape == (2, 3, 4)

# 袋 k たちを平均的なただ 1 つの袋 k にする
assert x.mean(dim=0, keepdim=True).shape == (1, 3, 4)
assert x.mean(dim=1, keepdim=True).shape == (2, 1, 4)
assert x.mean(dim=2, keepdim=True).shape == (2, 3, 1)

# 袋 k たちを平均的なただ 1 つの袋 k にした上で袋 k をやぶり捨てる (と袋 k + 1 が袋 k に昇格)
assert x.mean(dim=0, keepdim=False).shape == (3, 4)
assert x.mean(dim=1, keepdim=False).shape == (2, 4)
assert x.mean(dim=2, keepdim=False).shape == (2, 3)

x.mean(dim=(k, l), keepdim=True) は箱の中身を「平均的袋 l の平均的袋 k」にまとめ、x.mean(dim=None, keepdim=True) は箱の中身をただ 1 つの平均的だんごにまとめます。

import torch
x = torch.ones(2, 3, 4)
assert x.shape == (2, 3, 4)

assert x.mean(dim=(0, 1), keepdim=True).shape == (1, 1, 4)
assert x.mean(dim=None, keepdim=True).shape == (1, 1, 1)

x.unsqueeze

x.unsqueeze(dim=k) は新しい袋をもってきて袋 k たちをすっぽり包みます。k = d - 1 のときは最も内側のだんごたちをすっぽり包みます。k = d のときはだんごを個包装します。
x.unsqueeze(dim=k) は箱どうしをブロードキャスト可能にするときによく使います。

import torch

# 箱の中の 2 袋の 3 袋の 4 個入りだんご
x = torch.ones(2, 3, 4)
assert x.shape == (2, 3, 4)

# 袋 0 たちを包む (これが袋 0 に、元の袋 0 は 袋 1 に、元の袋 1 は 袋 2 になる)
assert x.unsqueeze(0).shape == (1, 2, 3, 4)

# 袋 1 たちを包む (これが袋 1 に、元の袋 1 は 袋 2 になる)
assert x.unsqueeze(1).shape == (2, 1, 3, 4)

# 袋 1 の内側のだんごたちを包む　(これが袋 2 になる)
assert x.unsqueeze(2).shape == (2, 3, 1, 4)

# だんごを個包装する　(これが袋 2 になる)
assert x.unsqueeze(3).shape == (2, 3, 4, 1)

ちなみに、x.mean(dim=k).unsqueeze(k) というコードをみることがありますが (袋 k を破り捨て、新しい袋 k たちを再度包む)、これなら x.mean(dim=k, keepdim=True) と同じです。

import torch

# 箱の中の 1 袋の 2 袋の 3 個入りだんご
x = torch.tensor([[[1., 2., 3.], [4., 5., 6.]]])
assert x.shape == (1, 2, 3)

# 袋 1 たちを平均的な一つの袋にし、袋 1 を破り捨て、新しい袋 1 たちを再度包む
x0 = x.mean(dim=1).unsqueeze(1)
assert x0.shape == (1, 1, 3)
assert torch.allclose(x0, torch.tensor([[[2.5, 3.5, 4.5]]]))

# 袋 1 たちを平均的な一つの袋にする
x1 = x.mean(dim=1, keepdim=True)
assert x1.shape == (1, 1, 3)
assert torch.allclose(x1, torch.tensor([[[2.5, 3.5, 4.5]]]))

ブロードキャスト可能

要素ごとの演算ができるテンソルどうしは「ブロードキャスト可能」といいます。同じ .shape をもつテンソルどうしは当然ブロードキャスト可能ですが、その他に以下の場合にもブロードキャスト可能です (＝共通の形状に拡張した後に要素ごとの演算がなされます)。

箱に入っていないだんご (0 階のテンソル) は、どんな箱ともブロードキャスト可能です (相手の箱の形に合わせてそのだんごを複製できるので)。
箱どうしで階数が同じ場合、袋 or だんご数がずれている階があっても片方の次元数が 1 ならブロードキャスト可能です (相手の箱の形に合わせて袋 or だんごを複製できるので)。
箱どうしで階数が異なる場合、階数が大きい方の .shape の右から階数が小さい方の階数だけをみて、階数が小さい方の箱とブロードキャスト可能ならブロードキャスト可能です (階数が小さい方の箱の中身を複製すればよいので)。