NumPyのブロードキャストを完全に理解する

1000点の観測点と2000点の候補地点があって、その全組み合わせ200万通りの距離を計算したい。Pythonの二重ループで書くと手元の環境で約1.9秒かかりますが、NumPyでは次の1行で、しかも実測で約90倍(約0.021秒)速く済みます。

import numpy as np

rng = np.random.default_rng(0)
X = rng.random((1000, 2))    # 観測点1000個の座標
Y = rng.random((2000, 2))    # 候補地点2000個の座標

D = np.sqrt(((X[:, None, :] - Y[None, :, :]) ** 2).sum(-1))   # ← この1行だけ
print(D.shape)   # (1000, 2000)

初めてこの行を見たとき、None が何をしているのか、なぜ形の違う配列同士を引き算できるのか、まったく分からなかったという方は多いはずです。ここで働いているのが ブロードキャスト(broadcasting) です。形の異なる配列を、NumPyが自動的に「揃えて」計算してくれる仕組みのことを指します。

ブロードキャストは、単なる書き方の省略ではありません。NumPyで速いコードを書くための中心的な道具です。

  • 機械学習では、全ペア距離行列やカーネル(グラム)行列 $K_{ij} = \exp(-\gamma\|\bm{x}_i – \bm{x}_j\|^2)$ の計算、特徴量ごとの標準化 $(\bm{A} – \bm{\mu})/\bm{\sigma}$、行方向のソフトマックス正規化が、すべてブロードキャスト1行で書けます。
  • 信号処理や物理シミュレーションでは、時刻の配列と周波数の配列から $e^{j2\pi f t}$ の行列を作ったり、格子点上のポテンシャル場を一気に評価したりする場面で、ブロードキャストなしには実用的な速度が出ません。

一方で、ブロードキャストは 静かに間違える 道具でもあります。形状 $(100,)$ の平均ベクトルを引くつもりが、うっかり $(100,1)$ の列ベクトルを引いてしまい、エラーも出ないまま $(100,100)$ の巨大な行列ができあがる。あるいは $N \times M \times D$ の中間配列がメモリを数十GB食いつぶしてプロセスが落ちる。こうした事故は、規則を「なんとなく」で使っている限り必ず起きます。

スカラーを配る・ベクトルを全行に配る・列ベクトルと行ベクトルから全組み合わせの表を作る、という3段階でブロードキャストの発想を示した概念図

この図が、記事全体で追いかけることの全体像です。左は誰もが無意識にやっているスカラー倍、真ん中は「長さ4のベクトルを3行すべてに配る」操作、右は「縦の値と横の値から全組み合わせの表を作る」操作で、3つとも同じ1つの規則で動いています。右へ行くほど強力になりますが、規則そのものは何も増えません。逆に言えば、左端が分かる人は右端も必ず分かるようになります。

この記事では、ブロードキャストの3つの規則を厳密に述べ、どの要素がどう複製されるのかを表と図で1ステップずつ追い、実務で頻出する3パターン(距離行列・正規化・外積グリッド)を実装し、最後にメモリの落とし穴とエラーメッセージの読み方までを扱います。読み終わったあとには、冒頭の1行がなぜ動くのかを、他人に説明できる状態になっているはずです。

本記事の内容

  • ブロードキャストの直感と、3つの規則の厳密な定式化
  • $(3,1)$ と $(1,4)$ が $(3,4)$ になるとき、どの要素がどこへ複製されるのか
  • None / np.newaxis / reshape / keepdims による意図的な次元追加
  • 内部実装(stride 0)と、それでもメモリを食う理由
  • 全ペア距離行列・正規化・外積グリッドの3パターン実装
  • ループ実装との速度比較、メモリ消費の実測、エラーメッセージの読み方

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

配列の shape、スライス、axis 引数の意味がおおよそ分かっていれば十分です。線形代数の知識は、最後の距離行列の展開式のところで少しだけ使います。

ブロードキャストとは — 「配る」という発想

表計算ソフトで、A列の全セルに同じ税率1.1を掛けたいとき、わざわざ1.1を全行に打ち込む人はいません。1つ書いてオートフィルすれば、同じ値が全行に「配られ」ます。ブロードキャストという言葉は、まさにこの「1つの値を多数へ配信する」イメージから来ています。放送(broadcast)が1つの電波を無数の受信機へ届けるのと同じです。

NumPyでもっとも身近なブロードキャストは、スカラーとの演算です。

import numpy as np

a = np.array([1.0, 2.0, 3.0])
print(a * 2)      # [2. 4. 6.]
print(a + 10)     # [11. 12. 13.]

出力を見ると、2 という1つの数値が3要素すべてに掛かっています。これは「$2$ を $[2, 2, 2]$ に展開してから要素ごとに掛けた」のと同じ結果です。誰もこれを不思議には思いませんが、実はここですでにブロードキャストが働いています。形の違うもの(スカラーと長さ3の配列)を、片方を仮想的に引き伸ばして揃えてから計算する — これがブロードキャストの本質です。

ブロードキャストは、この「スカラーの拡張」を多次元へ一般化したものだと考えてください。$(3,4)$ の行列に、長さ4のベクトルを足したいとします。意図としては「4列それぞれに対応するオフセットを、全3行に共通で足す」でしょう。

import numpy as np

A = np.zeros((3, 4))
offset = np.array([1.0, 2.0, 3.0, 4.0])   # 形状 (4,)
print(A + offset)

出力は3行とも [1. 2. 3. 4.] になります。長さ4のベクトルが、3行ぶんコピーされて足されたわけです。この「行方向へ配る」動作を、NumPyは形状の情報だけから自動的に決めています。ではその判断基準は何なのか。ここに一切の曖昧さはなく、たった3つの規則で完全に決まります。

重要なのは、ブロードキャストが 要素ごとの演算(ufunc)にだけ 適用されるという点です。+-*/**np.expnp.maximum、比較演算子 < などがこれに当たります。行列積 @np.dot は要素ごとの演算ではないので、別の規則(内側の次元が一致すること)に従います。ここを混同すると混乱の元なので、まず「ブロードキャスト=要素ごと演算のための形合わせ」と覚えておきましょう。

次のセクションで、その3つの規則を正確に述べます。規則そのものは驚くほど短いので、まずは条文として覚えてしまうのが近道です。

ブロードキャストの3つの規則

規則を述べる前に、なにを決めなければならないのかを整理しておきます。形状 $(3,4)$ の配列 $\bm{A}$ と形状 $(4,)$ の配列 $\bm{b}$ を足すとき、NumPyが決めるべきことは2つです。

  1. 結果の形状はどうなるか(ここでは $(3,4)$)
  2. 結果の $(i,j)$ 要素を作るのに、元の配列のどの要素を使うか(ここでは $A[i,j] + b[j]$)

この2つを、形状のタプルだけから機械的に決めるのが以下の規則です。全体の流れを先に図で示しておきます。

形状タプルを用意し、規則1で右揃え、規則2で軸ごとに判定、通れば伸長して結果の形状は軸ごとの最大値、通らなければ規則3でValueErrorになる判定フロー図

フローの入口から出口まで、判断材料が 形状タプルだけ であることに注目してください。配列の中身がどんな値でも、dtypeが何でも、結果の形状は同じです。だからこそ np.broadcast_shapes のように、実際の配列を1つも作らずに結果の形状を予言できる関数が成立します。分岐は1か所しかなく、そこで落ちれば ValueError、通れば伸長して軸ごとの最大値が結果になる、それだけの機械です。

規則1: 次元数を右揃えし、足りない側の左に 1 を補う

2つの配列の次元数(ndim)が違うとき、形状タプルを右端で揃え、次元数が少ない側の左側に長さ1の軸を必要なだけ補います。

$(3,4)$ と $(4,)$ なら、右揃えして

A: (3, 4)
b: (   4)   ->  (1, 4)  と見なす

となります。$(2,3,4)$ と $(3,4)$ なら $(1,3,4)$ と見なされます。

なぜ「左」に補うのでしょうか。NumPyの配列は、末尾の軸ほど「内側」(メモリ上で連続、行方向)を表す C オーダーが既定です。$(3,4)$ の行列は「3行 $\times$ 4列」であり、最後の軸が列に対応します。長さ4のベクトルは、直感的には「4列ぶんの値」と解釈するのが自然です。末尾の軸を意味的にそろえる という慣習が、右揃えという設計に反映されています。

規則2: 各軸について、長さが一致するか、どちらかが 1 なら通す

右揃えして次元数を揃えたら、軸ごとに独立に次の判定をします。

  • 2つの長さが等しい $\Rightarrow$ そのまま。結果の長さも同じ
  • どちらかが $1$ $\Rightarrow$ その軸を、もう一方の長さまで 引き伸ばす(同じ値を繰り返す)。結果の長さは大きいほう
  • 両方が $1$ $\Rightarrow$ 結果も $1$

引き伸ばしのことを stretch伸長 と呼びます。長さ1の軸は「まだ値が1つしかないが、いくらでも複製してよい軸」と解釈されるわけです。

規則3: 上記のどちらにも当てはまらない不一致はエラー

長さが $3$ と $4$ のように、一致もせず、どちらも $1$ でもない軸が1つでもあれば、その時点で ValueError になります。「長いほうに合わせて短いほうを繰り返す」ような親切な処理は 一切しません。$4$ を $8$ に伸ばす(2回繰り返す)ようなことも起きません。伸長が許されるのは 長さ1の軸だけ です。

規則の形式的なまとめ

数式で書いておきましょう。配列 $\bm{A}$ の形状を $(a_1, a_2, \dots, a_n)$、$\bm{B}$ の形状を $(b_1, b_2, \dots, b_m)$ とします。$r = \max(n, m)$ とおき、規則1に従って左側に1を補った形状を

$$ a’_k = \begin{cases} 1 & (k \le r – n) \\ a_{k – (r-n)} & (k > r – n) \end{cases} \qquad b’_k = \begin{cases} 1 & (k \le r – m) \\ b_{k – (r-m)} & (k > r – m) \end{cases} $$

と定義します。ここで $k = 1, \dots, r$ です。すると、ブロードキャスト可能である条件は

$$ \forall k: \quad a’_k = b’_k \ \ \text{または} \ \ a’_k = 1 \ \ \text{または} \ \ b’_k = 1 $$

であり、このとき結果の形状 $(c_1, \dots, c_r)$ は

$$ c_k = \max(a’_k,\, b’_k) $$

で与えられます。ここまでが「規則1と2と3」の内容です。

さらに、結果のどの要素が元のどの要素から来るのか も式で書けます。インデックスの写像を

$$ \iota(i; s) = \begin{cases} i & (s > 1) \\ 0 & (s = 1) \end{cases} $$

と定義すると、結果 $\bm{C}$ の要素は

$$ C[i_1, \dots, i_r] = f\big(A[\iota(i_1; a’_1), \dots, \iota(i_r; a’_r)],\ B[\iota(i_1; b’_1), \dots, \iota(i_r; b’_r)]\big) $$

となります($f$ は + などの2項演算)。式が仰々しく見えますが、言っていることは単純です。その軸の長さが1なら、どんなインデックスで参照されても常に0番目の要素を返す。これが「引き伸ばし」の正体です。値をコピーしているのではなく、同じ要素を何度も読みに行っているだけ なのです。この観察は、あとでメモリの話をするときに効いてきます。

具体例で規則を回してみる

いくつかの組み合わせを、規則に沿って手で判定してみましょう。

$\bm{A}$ の形状 $\bm{B}$ の形状 右揃え後 判定 結果
$(3,4)$ $(4,)$ $(3,4)$ vs $(1,4)$ 軸0: 3 vs 1 → OK / 軸1: 4 vs 4 → OK $(3,4)$
$(3,1)$ $(1,4)$ $(3,1)$ vs $(1,4)$ 軸0: 3 vs 1 → OK / 軸1: 1 vs 4 → OK $(3,4)$
$(3,4)$ $(3,)$ $(3,4)$ vs $(1,3)$ 軸1: 4 vs 3 → NG エラー
$(3,4)$ $(3,1)$ $(3,4)$ vs $(3,1)$ 軸0: 一致 / 軸1: 4 vs 1 → OK $(3,4)$
$(5,1,3)$ $(1,4,1)$ そのまま 全軸OK $(5,4,3)$
$(2,3,4)$ $(3,)$ $(2,3,4)$ vs $(1,1,3)$ 軸2: 4 vs 3 → NG エラー
$(8,1,6,1)$ $(7,1,5)$ $(8,1,6,1)$ vs $(1,7,1,5)$ 全軸OK $(8,7,6,5)$

この表を、実際に右揃えした状態で見比べられるように図にしておきます。

7つの形状の組み合わせを右端でそろえ、軸ごとに緑(通る)と赤(一致せず1でもない)で色分けし、結果の形状またはValueErrorを示した判定図

白い破線のマスが、規則1によって「勝手に補われた長さ1」です。エラーになる2ケース(3段目と6段目)はどちらも、いちばん右の軸で赤が出ています。エラーは必ず特定の1軸で起きるので、赤いマスを1つ見つけた時点で原因の特定は終わりです。最下段の $(8,1,6,1)$ と $(7,1,5)$ では、上の配列が軸1で、下の配列が軸0・軸2でそれぞれ伸長され、両者が別々の軸で歩み寄っている様子が読み取れます。

3行目の $(3,4)$ と $(3,)$ がエラーになるのは、初学者がもっともつまずくところです。「3行に対して各行の値を足したい」つもりでも、$(3,)$ は右揃えで $(1,3)$ と解釈され、列方向の長さ3 として扱われてしまいます。行方向に配りたいなら、自分で $(3,1)$ に整形しなければなりません。この「自動補完は必ず左にしか入らない」という非対称性が、次のセクションで None を使う理由そのものになります。

最下段の $(8,1,6,1)$ と $(7,1,5)$ は、公式ドキュメントにも載っている典型例です。両方の配列がそれぞれ別の軸で伸長されて、4次元の結果になります。ブロードキャストが「片方を大きいほうへ合わせる」ものではなく、軸ごとに独立に、大きいほうへ双方が歩み寄る ものだと分かる例です。

規則は以上ですべてです。しかし「規則を知っている」と「どの要素がどこへ行くかが見える」の間には、まだ距離があります。次のセクションで、いちばん基本的な $(3,1) + (1,4)$ を、要素レベルで完全に追跡します。

$(3,1)$ と $(1,4)$ が $(3,4)$ になる仕組みを要素で追う

ここが本記事の心臓部です。もっとも典型的なブロードキャスト、列ベクトル $+$ 行ベクトルを、要素レベルで完全に追跡します。

$\bm{a}$ を形状 $(3,1)$ の列ベクトル、$\bm{b}$ を形状 $(1,4)$ の行ベクトルとします。

$$ \bm{a} = \begin{pmatrix} 0 \\ 1 \\ 2 \end{pmatrix}, \qquad \bm{b} = \begin{pmatrix} 0 & 1 & 2 & 3 \end{pmatrix} $$

規則1では、どちらも2次元なので補う必要はありません。規則2で軸ごとに見ると、軸0が $3$ vs $1$ なので $\bm{b}$ 側を3行に伸長、軸1が $1$ vs $4$ なので $\bm{a}$ 側を4列に伸長します。結果の形状は $(\max(3,1), \max(1,4)) = (3,4)$ です。

伸長後の中身を書き下しましょう。$\bm{a}$ は軸1の長さが1なので、$\iota(j;1) = 0$、つまり列インデックス $j$ が何であっても $a[i,0]$ を読みます。

$$ \tilde{\bm{a}} = \begin{pmatrix} 0 & 0 & 0 & 0 \\ 1 & 1 & 1 & 1 \\ 2 & 2 & 2 & 2 \end{pmatrix} $$

同様に $\bm{b}$ は軸0の長さが1なので、行インデックス $i$ が何であっても $b[0,j]$ を読みます。

$$ \tilde{\bm{b}} = \begin{pmatrix} 0 & 1 & 2 & 3 \\ 0 & 1 & 2 & 3 \\ 0 & 1 & 2 & 3 \end{pmatrix} $$

この2つを要素ごとに足すので、結果は

$$ \tilde{\bm{a}} + \tilde{\bm{b}} = \begin{pmatrix} 0 & 1 & 2 & 3 \\ 1 & 2 & 3 & 4 \\ 2 & 3 & 4 & 5 \end{pmatrix} $$

となります。一般式で書けば

$$ C[i, j] = a[i, 0] + b[0, j] = a_i + b_j $$

です。つまり $(3,1) + (1,4)$ は、すべての $(i,j)$ の組み合わせについて $a_i + b_j$ を並べた表 を作る操作にほかなりません。これが「外積的な演算」の正体です。掛け算にすれば $C[i,j] = a_i b_j$ で、これはまさにベクトルの外積(テンソル積)$\bm{a}\bm{b}^\top$ です。

この一連の流れを、5枚のパネルで一気に見てしまいましょう。

形状(3,1)のaと形状(1,4)のb、それぞれを(3,4)に伸長した様子、そして両者の和を5枚のヒートマップで並べ、伸長で増えたセルを赤い破線で示した図

赤い破線で囲まれたマスが「伸長によって増えた」セルです。a は横方向に同じ値が4回、b は縦方向に同じ行が3回並び、どちらも もとの配列には存在しない値が1つも増えていない ことに注目してください。右端の和は $C[i,j] = a_i + b_j$、つまり左上の0から右下の5まで、行と列の番号を足した表そのものになっています。

どの要素がどこから来るかを表にまとめておきます。

結果の位置 $(i,j)$ $\bm{a}$ の参照先 $\bm{b}$ の参照先
$(0,0)$ $a[0,0]=0$ $b[0,0]=0$ $0$
$(0,3)$ $a[0,0]=0$ $b[0,3]=3$ $3$
$(2,0)$ $a[2,0]=2$ $b[0,0]=0$ $2$
$(2,3)$ $a[2,0]=2$ $b[0,3]=3$ $5$
$(1,2)$ $a[1,0]=1$ $b[0,2]=2$ $3$

表を眺めると、$\bm{a}$ の参照先は 列インデックス $j$ に依存していない、$\bm{b}$ の参照先は 行インデックス $i$ に依存していない ことがはっきりします。「長さ1の軸は、そのインデックスを無視する軸」だと理解すると、複雑な多次元ブロードキャストも一気に見通しがよくなります。

コードで確認しておきましょう。

import numpy as np

a = np.arange(3).reshape(3, 1)   # 形状 (3, 1)
b = np.arange(4).reshape(1, 4)   # 形状 (1, 4)

print("a.shape =", a.shape, " b.shape =", b.shape)
print("結果の形状:", np.broadcast_shapes(a.shape, b.shape))
print(a + b)

実行すると 結果の形状: (3, 4) と表示され、続いて上で手計算した行列がそのまま出力されます。手で追った $\tilde{\bm{a}} + \tilde{\bm{b}}$ と一致することを確認してください。np.broadcast_shapes は実際の配列を作らずに 形状だけを予測 してくれる関数で、デバッグの強い味方になります。大きな配列で試す前に、まずこれで形の見積もりをするのが安全です。

ここで、$(3,1)+(1,4)$ と $(3,)+(4,)$ の違いにも注意しておきましょう。後者は右揃えで $(3,)$ vs $(4,)$ となり、軸が1つしかないので $3 \ne 4$ でエラーです。「全組み合わせを作りたい」なら、片方を列ベクトルに、もう片方を行ベクトルに整形する必要がある — これが次のセクションのテーマです。

意図的に次元を足す — None / np.newaxis / reshape

前のセクションの最後で見たとおり、$(3,)$ と $(4,)$ をそのまま足すことはできません。「全組み合わせの表を作る」には、$(3,1)$ と $(4,)$ にする必要があります(後者は自動で $(1,4)$ と解釈されます)。この整形を簡潔に書くのが None です。

import numpy as np

u = np.arange(1, 4)   # 形状 (3,)
v = np.arange(1, 5)   # 形状 (4,)

print(u[:, None].shape)   # (3, 1)  末尾に長さ1の軸を追加
print(v[None, :].shape)   # (1, 4)  先頭に長さ1の軸を追加
print(u[:, None] * v[None, :])

出力は形状 $(3,4)$ の九九表のような行列 [[1,2,3,4],[2,4,6,8],[3,6,9,12]] になります。$C[i,j] = u_i v_j$ という定義どおりの外積が得られました。u[:, None]None はスライスの中に置かれた「新しい軸をここに挿入せよ」という指示で、np.newaxis はこの None の別名にすぎません(実際に np.newaxis is NoneTrue を返します)。読みやすさのために np.newaxis と書く流儀もありますが、意味は完全に同じです。

4つの書き方の比較

次元を足す方法は複数あります。使い分けを整理しておきましょう。

書き方 特徴
x[:, None] $(3,) \to (3,1)$ 最短。挿入位置がスライス中の位置で直感的に分かる
x[None, :] $(3,) \to (1,3)$ 先頭に追加。実は右揃えで自動的にこうなるので省略可
x.reshape(3, 1) $(3,) \to (3,1)$ 明示的だが、要素数を書くのでコードが硬くなる
x.reshape(-1, 1) $(n,) \to (n,1)$ 長さ可変でも動く。None と並んでよく使う
np.expand_dims(x, 1) 軸1に挿入 軸番号を変数で渡せる。汎用コードで便利
.sum(axis=1, keepdims=True) $(3,4) \to (3,1)$ 集約と同時に軸を潰さず残す。正規化で必須

もっとも重要なのは、「先頭への追加は自動、末尾への追加は手動」という非対称性です。 規則1が「左に1を補う」と決めている以上、$(3,) \to (1,3)$ は放っておいてもNumPyがやってくれます。逆に $(3,) \to (3,1)$ は絶対に自動では起きません。したがって実務では、

$$ \underbrace{u[:, \text{None}]}_{(3,1)} \ \ast \ \underbrace{v}_{(4,) \to (1,4)} $$

のように 片方だけに [:, None] を書けば十分 です。両方に書く u[:, None] * v[None, :] は冗長ですが、「これは全組み合わせを作る演算だ」という意図が読み手に伝わる ので、私はあえて両方書くことを勧めます。半年後の自分がコードを読むときの可読性は、3文字ぶんのタイプより価値があります。

この非対称性は、図で並べて見るともう忘れません。

左は(3,4)と長さ4のベクトルが自動で通る様子、右は(3,4)と長さ3のベクトルが軸1で4対3となりValueErrorになり、[:, None]で(3,1)にすれば通ることを示した比較図

左右の図は、やりたいことが「列ごとに配る」か「行ごとに配る」かが違うだけで、配列の大きさも値も同じです。それでも左は何も書かずに通り、右は [:, None] を書かないと通りません。この差を生んでいるのは規則1の「左に補う」という一文だけ です。右図で長さ3のベクトルが $(1,3)$ と読まれ、軸1で 4 対 3 となって落ちる様子を、赤いバツ印で示してあります。

keepdims の役割

集約関数(sum, mean, max, std など)は既定でその軸を消します。$(3,4)$ を axis=1 で合計すると $(3,)$ になり、$(3,)$ は右揃えで $(1,3)$ と解釈されるため、元の $(3,4)$ とはブロードキャストできません。ここで keepdims=True を付けると、軸を長さ1で残して $(3,1)$ になり、$(3,4)$ とそのままブロードキャストできます。

import numpy as np

A = np.arange(12, dtype=float).reshape(3, 4)
print(A.sum(axis=1).shape)                  # (3,)   軸が消える
print(A.sum(axis=1, keepdims=True).shape)   # (3, 1) 軸が残る

# 各行を「行の合計」で割る(行方向の正規化)
P = A / A.sum(axis=1, keepdims=True)
print(P.sum(axis=1))    # [1. 1. 1.]

出力の [1. 1. 1.] は、各行の合計がちょうど1に正規化されたことを示しています。もし keepdims=True を忘れると、$(3,4)$ を $(3,)$ で割ろうとして「軸1が 4 vs 3」でエラーになります。keepdims=True は「あとでブロードキャストで元の配列に戻す」ための予約札 だと覚えておきましょう。これを付けるかどうかの判断は、「集約結果を元の配列と組み合わせて使うか?」で決まります。使うなら付ける、表示するだけなら要りません。

次は、この「引き伸ばし」がメモリ上で実際に何をしているのかを覗いてみます。ここを知ると、ブロードキャストが速い理由と、それでもメモリを食う理由の両方が同時に理解できます。

内部実装: 伸長は stride 0 で実現されている

「長さ1の軸を伸ばす」と聞くと、値を物理的にコピーして大きな配列を作っているように思えます。しかしNumPyはコピーしません。代わりに ストライド(stride)を0にする という巧妙な手を使います。

ストライドとは、「その軸のインデックスを1増やすとき、メモリ上で何バイト進むか」を表す数値です。float64(8バイト)の $(3,4)$ 配列なら、ストライドは $(32, 8)$ です。列を1つ進むと8バイト、行を1つ進むと $4 \times 8 = 32$ バイト進みます。

ここで、$(3,1)$ の配列を $(3,4)$ に伸長した「ビュー」を作ってみましょう。

import numpy as np

a = np.arange(3).reshape(3, 1)         # int64, 形状 (3,1)
bt = np.broadcast_to(a, (3, 4))        # (3,4) に伸長したビュー

print("元のストライド  :", a.strides)   # (8, 8)
print("伸長後ストライド:", bt.strides)  # (8, 0)  ← 列方向が 0
print("メモリ共有      :", np.shares_memory(bt, a))   # True
print("見かけの nbytes :", bt.nbytes, " 実体:", a.nbytes)
print("書き込み可能    :", bt.flags.writeable)        # False

実行結果は次のようになります。伸長後のストライドは (8, 0)、メモリ共有は Truenbytes は96と表示されるのに実体は24バイト、そして書き込みは禁止されています。

起きていることを絵にすると次のようになります。

左に24バイトの実体(0,1,2の3要素)、右に見かけ上の(3,4)配列を描き、1行目の4マスがすべて同じ8バイトを指す矢印とstrides=(8,0)などの実測値を示した内部実装の模式図

右の $(3,4)$ に見える配列は、実際には左の24バイトしか持っていません。橙色の矢印が示すとおり、1行目の4マスは すべて同じ8バイトを指しています。破線のマスは「表示されるが実体を持たないセル」で、nbytes が96と答えるのはこの破線ぶんまで数えてしまうからです。書き込みが禁止されている理由も図から明らかで、破線のどれか1つに書けば実線のマスごと4か所が同時に変わってしまいます。

この 0 がすべてを説明します。列インデックスを1増やしても、メモリ上のポインタは1バイトも動かない。 つまり bt[0,0]bt[0,1]bt[0,2]bt[0,3] はすべて 同一の8バイト を指しています。値を4個コピーしたのではなく、同じ1個を4回読んでいるだけです。前のセクションで見たインデックス写像 $\iota(j;1)=0$ が、そのままストライド0として実装されているわけです。

3つの重要な帰結があります。

1つ目。nbytes は嘘をつきます。 bt.nbytes が96と出るのは、nbytes が「要素数 $\times$ 要素サイズ」を単純計算しているからです。実際に確保されているのは元の24バイトだけです。ビューのメモリ実消費を知りたいときは nbytes ではなく、base をたどるか np.shares_memory で確認する必要があります。

2つ目。broadcast_to の結果は読み込み専用です。 書き込みを許すと、bt[0,1] = 99 としたつもりが bt[0,0] から bt[0,3] まで全部99になってしまいます。同じメモリを指しているのだから当然です。この事故を防ぐため、NumPyは書き込み禁止フラグを立てます。どうしても書き込みたければ .copy() して実体化するしかありません。

3つ目、そしてもっとも大事な点。ビューのままで済むのは「入力側」だけです。 a + b を計算した瞬間、その 結果は必ず新しい実体として確保されます。$(3,1) + (1,4)$ なら、入力は3要素と4要素しか持っていなくても、出力は12要素ぶんのメモリを新規に確保します。$(10000,1) + (1,10000)$ なら、入力は合計16万バイト程度なのに、出力は $10^8 \times 8 = 800$ MB です。「ブロードキャストはコピーしないから安心」は半分しか正しくない のです。これがのちほど扱うメモリの落とし穴の正体です。

なお、ストライド0の副産物として、ブロードキャストは キャッシュ効率がよい という利点もあります。同じ8バイトを何度も読むだけなので、その値はCPUのL1キャッシュに載り続けます。Pythonの二重ループが遅いのは、ループ1周ごとにPythonオブジェクトの生成・型チェック・関数呼び出しが走るからで、ブロードキャストはこれをすべてC言語のループに置き換えます。速度差が2桁になるのはこのためです。

仕組みが分かったところで、実務で頻出する3つのパターンを見ていきましょう。まずは冒頭で予告した全ペア距離行列です。

パターン1: 全ペア距離行列を1ステップずつ追う

$N$ 個の点 $\bm{x}_1, \dots, \bm{x}_N \in \mathbb{R}^D$ と、$M$ 個の点 $\bm{y}_1, \dots, \bm{y}_M \in \mathbb{R}^D$ があるとき、すべてのペアのユークリッド距離

$$ d_{ij} = \|\bm{x}_i – \bm{y}_j\| = \sqrt{\sum_{k=1}^{D} (x_{ik} – y_{jk})^2} $$

を並べた $N \times M$ 行列を作りたい、という場面はk近傍法・クラスタリング・カーネル法・粒子シミュレーションなど至るところに現れます。

素直に書けば二重ループですが、ブロードキャストなら1行です。ただしその1行は、形状の変化を追わないと絶対に理解できません。1ステップずつ分解します。

ステップ1: 3次元へ持ち上げる

$\bm{X}$ の形状は $(N, D)$、$\bm{Y}$ の形状は $(M, D)$ です。目標の $d_{ij}$ は $i$ と $j$ の2つの添字を持ち、さらに $k$ で和を取るので、途中で $(N, M, D)$ の3階テンソルが必要 になります。そこで各配列に軸を1つ挿入します。

$$ \bm{X}[:, \text{None}, :] : (N, 1, D), \qquad \bm{Y}[\text{None}, :, :] : (1, M, D) $$

$\bm{X}$ には「$j$ に対応する軸」を、$\bm{Y}$ には「$i$ に対応する軸」を、長さ1で挿入したことになります。長さ1の軸は「そのインデックスを無視する軸」でした。つまり $\bm{X}$ 側は $j$ を無視し、$\bm{Y}$ 側は $i$ を無視する — これはまさに私たちが望む挙動です。

ステップ2: 引き算でブロードキャストが起きる

軸ごとに規則2を適用します。

$\bm{X}[:,None,:]$ $\bm{Y}[None,:,:]$ 結果
0 $N$ $1$ $N$($\bm{Y}$ を $N$ 回伸長)
1 $1$ $M$ $M$($\bm{X}$ を $M$ 回伸長)
2 $D$ $D$ $D$(一致)

結果は $(N, M, D)$ で、その中身は

$$ \Delta[i, j, k] = X[i, k] – Y[j, k] = x_{ik} – y_{jk} $$

です。$\Delta[i,j,:]$ が、ちょうど差ベクトル $\bm{x}_i – \bm{y}_j$ になっています。

ステップ3: 二乗して最終軸で和を取り、平方根

$$ d_{ij} = \sqrt{\sum_{k} \Delta[i,j,k]^2} $$

なので、(diff ** 2).sum(axis=-1) で $(N, M)$ に落とし、np.sqrt を掛ければ完成です。axis=-1 は「最後の軸」を意味し、$D$ の値に依存しないので安全です。

形状の変遷をまとめると次のようになります。

$$ (N,D) \ \& \ (M,D) \ \xrightarrow{\text{None挿入}} \ (N,1,D) \ \& \ (1,M,D) \ \xrightarrow{\ -\ } \ (N,M,D) \ \xrightarrow{\ \text{sum}(-1)\ } \ (N,M) $$

(N,D)と(M,D)からNoneで軸を挿入して(N,1,D)と(1,M,D)にし、引き算で(N,M,D)、最終軸の二乗和と平方根で(N,M)になる4段階の形状変遷図

4つの箱の中で、メモリを大量に食うのは3つ目の $(N,M,D)$ だけ です。入力も、None を挿入した2つ目も、ストライド0のビューなので追加のメモリはほぼゼロ。ところが引き算をした瞬間に $N M D \times 8$ バイトが実体として確保されます。この図の3つ目の箱を見た瞬間に「何バイトになるか」を暗算できるようになると、事故はほぼ起きなくなります。

import numpy as np

rng = np.random.default_rng(0)
X = rng.normal(size=(5, 2))   # 5点、2次元
Y = rng.normal(size=(7, 2))   # 7点、2次元

diff = X[:, None, :] - Y[None, :, :]
print("X[:, None, :] の形状:", X[:, None, :].shape)   # (5, 1, 2)
print("Y[None, :, :] の形状:", Y[None, :, :].shape)   # (1, 7, 2)
print("差の形状            :", diff.shape)             # (5, 7, 2)

D_mat = np.sqrt((diff ** 2).sum(axis=-1))
print("距離行列の形状      :", D_mat.shape)            # (5, 7)
print("D[2,3] =", D_mat[2, 3], " 直接計算 =", np.linalg.norm(X[2] - Y[3]))

出力では D[2,3]np.linalg.norm(X[2]-Y[3]) が小数点以下まで完全に一致します。形状も $(5,1,2) \to (5,7,2) \to (5,7)$ と、設計どおりに変遷していることが確認できます。ここまで追えば、冒頭の1行はもう暗号ではなくなったはずです。

できあがった距離行列が、本当に「全ペアの距離」になっているかを目で確かめておきましょう。

左に6個のX点と8個のY点を結ぶ48本の線を描いた散布図、右に対応する形状(6,8)の距離行列ヒートマップを並べ、最短ペアと最長ペアを緑と赤で対応づけた図

左の灰色の線1本1本が、右のヒートマップの1マスに対応しています。緑で示した最短ペア $d = 1.47$ は右図の $(i,j) = (5,1)$ の暗いマス、赤の最長ペア $d = 4.87$ は $(3,7)$ の明るいマスにちょうど対応しています。行が $\bm{X}$ の番号、列が $\bm{Y}$ の番号 という並びも、$(N,M)$ という形状そのままです。左図で近い点どうしほど右図が暗くなる、という対応が取れていれば、距離行列は正しく作れています。

展開式による省メモリ版

$(N,M,D)$ の中間配列は、$D$ が大きいと重くのしかかります。そこで、二乗距離を展開してみましょう。

$$ d_{ij}^2 = \|\bm{x}_i – \bm{y}_j\|^2 = (\bm{x}_i – \bm{y}_j)^\top (\bm{x}_i – \bm{y}_j) $$

右辺を展開すると、

$$ d_{ij}^2 = \bm{x}_i^\top \bm{x}_i – 2\,\bm{x}_i^\top \bm{y}_j + \bm{y}_j^\top \bm{y}_j $$

となります。ここで $\bm{x}_i^\top\bm{x}_i = \|\bm{x}_i\|^2$ は $i$ だけの関数、$\|\bm{y}_j\|^2$ は $j$ だけの関数、そして $\bm{x}_i^\top\bm{y}_j$ は行列積 $\bm{X}\bm{Y}^\top$ の $(i,j)$ 成分です。したがって、$\bm{s} = (\|\bm{x}_i\|^2)_i \in \mathbb{R}^N$、$\bm{t} = (\|\bm{y}_j\|^2)_j \in \mathbb{R}^M$ とおけば

$$ \bm{D}^{(2)} = \underbrace{\bm{s}[:, \text{None}]}_{(N,1)} + \underbrace{\bm{t}[\text{None}, :]}_{(1,M)} – 2\,\underbrace{\bm{X}\bm{Y}^\top}_{(N,M)} $$

と書けます。ここでも $(N,1) + (1,M) \to (N,M)$ というブロードキャストが登場しました。中間で $(N,M,D)$ を一切作らずに済む のがポイントです。$D$ が100や1000のときには、この差が決定的になります。

import numpy as np

def pairwise_dist_gemm(X, Y):
    """展開式による距離行列。中間で (N,M,D) を作らない"""
    s = (X ** 2).sum(axis=1)          # (N,)  各行のノルム二乗
    t = (Y ** 2).sum(axis=1)          # (M,)
    d2 = s[:, None] + t[None, :] - 2.0 * (X @ Y.T)   # (N,M)
    return np.sqrt(np.maximum(d2, 0.0))   # 丸め誤差で負になるのを防ぐ

rng = np.random.default_rng(0)
X = rng.normal(size=(200, 8)); Y = rng.normal(size=(300, 8))
D1 = np.sqrt(((X[:, None, :] - Y[None, :, :]) ** 2).sum(-1))
D2 = pairwise_dist_gemm(X, Y)
print("最大絶対誤差:", np.abs(D1 - D2).max())

最大絶対誤差は実測で $1.8 \times 10^{-15}$、つまり倍精度の丸め誤差の範囲に収まります。2つの式が数学的に等価であることが数値的にも確認できました。np.maximum(d2, 0.0) を挟んでいるのは、同一点同士のように理論値が0の場合に、浮動小数の引き算で $-10^{-16}$ のような微小な負値が出ることがあり、そのまま sqrt に渡すと nan になるからです。展開式を使うときの定番の防御策なので、必ず入れておきましょう。

距離行列が書けるようになると、カーネル法もすぐそこです。RBFカーネル $K_{ij} = \exp(-\gamma d_{ij}^2)$ は、上で作った d2np.exp(-gamma * d2) を適用するだけです。次のパターンでは、もっと日常的に使う「正規化」を見ていきます。

パターン2: 正規化・標準化と (n,) vs (n,1) の罠

データ分析でもっとも頻繁に書くブロードキャストは、おそらく標準化(z-score化)でしょう。形状 $(N, D)$ のデータ行列 $\bm{A}$ について、特徴量ごとに 平均0・標準偏差1へそろえます。

$$ Z_{ij} = \frac{A_{ij} – \mu_j}{\sigma_j}, \qquad \mu_j = \frac{1}{N}\sum_{i=1}^{N} A_{ij}, \qquad \sigma_j = \sqrt{\frac{1}{N}\sum_{i=1}^{N}(A_{ij}-\mu_j)^2} $$

$\mu_j$ と $\sigma_j$ は列ごとの量なので、A.mean(axis=0) で形状 $(D,)$ の配列として得られます。これを $(N,D)$ から引くと、右揃えで $(1,D)$ と解釈され、全行に同じ平均ベクトルが配られます。まさに欲しい動作です。

import numpy as np

rng = np.random.default_rng(42)
# 列ごとに平均も分散も違うデータを作る
A = rng.normal(loc=[0, 10, -5], scale=[1, 3, 0.5], size=(1000, 3))

mu = A.mean(axis=0)   # (3,)
sd = A.std(axis=0)    # (3,)
print("列ごとの平均:", np.round(mu, 3))
print("列ごとの標準偏差:", np.round(sd, 3))

Z = (A - mu) / sd     # (1000,3) - (3,) -> ブロードキャストで (1000,3)
print("標準化後の平均:", np.round(Z.mean(axis=0), 12))
print("標準化後の標準偏差:", np.round(Z.std(axis=0), 12))

出力を見ると、元データの平均は約 [-0.008, 9.908, -5.019]、標準偏差は約 [1.013, 2.994, 0.505] と、指定した locscale にきちんと近い値になっています。この1行が何をしたのかを、分布の形で見ておきましょう。

左に平均0・10・-5、幅1・3・0.5の3つの特徴量のヒストグラムが離れて並ぶ標準化前の分布、右に3本が原点付近で重なる標準化後の分布を並べた図

左図では3本のヒストグラムが別々の場所に別々の幅で立っており、特徴量3(緑)だけ幅が極端に狭いことも分かります。右図ではその3本が原点に重なり、広がりもそろいました。図の下に添えた実測値のとおり、標準化後の平均は3列とも $10^{-14}$ 未満($1.5\times10^{-17}$、$1.1\times10^{-15}$、$-7.9\times10^{-15}$)、標準偏差は3列とも 1.000 です。列ごとに違う平均・標準偏差を、1000行すべてに正しく配った 結果がこの重なりです。標準化後は平均がすべて0(表示上は $\pm 0$)、標準偏差がすべて1になり、列ごとにスケールがそろった ことが確認できます。ここで axis=0 は「行方向に潰す=列ごとの統計量」を意味します。axis は「潰す軸」であって「残す軸」ではない、という点は何度でも確認しておきたいところです。

行方向の正規化では keepdims が要る

一方、各行のノルムを1にそろえる(コサイン類似度の前処理などで使う)場合は事情が変わります。

$$ U_{ij} = \frac{V_{ij}}{\|\bm{v}_i\|}, \qquad \|\bm{v}_i\| = \sqrt{\sum_{j} V_{ij}^2} $$

$\|\bm{v}_i\|$ は行ごとの量なので np.linalg.norm(V, axis=1) で形状 $(N,)$ になります。これで $(N,D)$ を割ろうとすると、右揃えで $(1,N)$ と解釈され、列方向の長さ $N$ として扱われてしまいます。$D \ne N$ ならエラー、そして恐ろしいことに $D = N$(正方行列)なら静かに間違った答えを返します

import numpy as np

rng = np.random.default_rng(1)
V = rng.normal(size=(3, 3))          # わざと正方行列
n = np.linalg.norm(V, axis=1)        # (3,)

W_wrong = V / n                      # エラーにならない! 列ごとに割ってしまう
W_ok    = V / n[:, None]             # (3,3) / (3,1) -> 正しく行ごと

print("誤り版の各行ノルム:", np.round(np.linalg.norm(W_wrong, axis=1), 4))
print("正解版の各行ノルム:", np.round(np.linalg.norm(W_ok,    axis=1), 4))

出力は、誤り版が [0.7242, 1.5632, 0.7864]、正解版が [1., 1., 1.] となります。何がずれているのかを図で確かめます。

左に(3,)が(1,3)と読まれ列ごとに割られる誤り版、中央に(3,1)で行ごとに割られる正解版、右に正規化後の行ノルムを比較した棒グラフ(誤り版0.724・1.563・0.786に対し正解版はすべて1.000)を並べた図

左と中央の格子は、割り算に使われる値がどう配られたかを示しています。誤り版では同じ値が 縦に 並び(=列ごとに同じ数で割る)、正解版では 横に 並んでいます(=行ごとに同じ数で割る)。右の棒グラフが結果で、誤り版の行ノルムは 0.724・1.563・0.786 とばらばら、正解版は3行とも 1.000 ちょうどです。この0.7や1.5という値は、例外も警告もいっさい出さずに黙って返ってきます — 行ノルムを実際に測りに行かない限り、間違いに気づく手がかりはどこにもありません。これが「静かに間違える」典型例です。$3 \times 4$ の非正方行列なら operands could not be broadcast together with shapes (4,3) (4,) とエラーが出て気づけるのに、正方行列だと通ってしまう。デバッグが困難なバグの温床です。

したがって、行方向に配りたいときは必ず keepdims=True[:, None] を書く ことを習慣にしてください。np.linalg.norm(V, axis=1, keepdims=True) と書けば最初から $(3,1)$ が得られ、この事故は起きません。列方向(axis=0)なら自動で正しく配られるので不要 — この非対称性は、規則1の「左に補う」から必然的に生じるものです。

数値的に安全なソフトマックス

ブロードキャストの合わせ技として、行ごとのソフトマックスも見ておきましょう。

$$ p_{ij} = \frac{\exp(z_{ij})}{\sum_{j’} \exp(z_{ij’})} = \frac{\exp(z_{ij} – m_i)}{\sum_{j’} \exp(z_{ij’} – m_i)}, \qquad m_i = \max_j z_{ij} $$

2つ目の等号は、分子分母に $\exp(-m_i)$ を掛けただけで、値は変わりません。しかし $z_{ij}$ が大きいと $\exp$ が容易にオーバーフローするため、行ごとの最大値 $m_i$ を引いてから指数を取る のが鉄則です。この $m_i$ の引き算と、分母での割り算が、どちらも keepdims=True によるブロードキャストになります。

import numpy as np

def softmax_rows(Z):
    """行ごとのソフトマックス(数値安定版)"""
    Z = np.asarray(Z, dtype=float)
    E = np.exp(Z - Z.max(axis=1, keepdims=True))   # (N,C) - (N,1)
    return E / E.sum(axis=1, keepdims=True)        # (N,C) / (N,1)

Z = np.array([[1.0, 2.0, 3.0],
              [1.0, 1.0, 1.0]])
P = softmax_rows(Z)
print(np.round(P, 4))
print("各行の和:", P.sum(axis=1))

出力は [[0.09, 0.2447, 0.6652], [0.3333, 0.3333, 0.3333]] で、各行の和はきっちり1です。1行目は値が大きいほど確率が高く、2行目は全要素が等しいので一様分布になる — どちらも定義から予想されるとおりの振る舞いです。この関数の中に for ループが1つもないこと、そして $(N,C)$ と $(N,1)$ の組み合わせが2回登場していることを確認してください。

正規化まで押さえたら、残るは格子・グリッドの生成です。ここで np.meshgrid との関係がはっきりします。

パターン3: 外積・格子生成と meshgrid との使い分け

2変数関数 $f(x,y)$ を格子上で評価して等高線を描きたい、という場面を考えます。$x$ 座標が $n_x$ 個、$y$ 座標が $n_y$ 個あるとき、必要なのは全組み合わせ $n_x \times n_y$ 点での値です。

もっとも素朴なのは np.meshgrid で座標行列を作る方法です。

import numpy as np

x = np.linspace(-2, 2, 5)   # (5,)
y = np.linspace(-1, 1, 3)   # (3,)

XX, YY = np.meshgrid(x, y)             # 既定は indexing='xy'
print("xy:", XX.shape, YY.shape)       # (3, 5) (3, 5)

XX2, YY2 = np.meshgrid(x, y, indexing='ij')
print("ij:", XX2.shape, YY2.shape)     # (5, 3) (5, 3)

XX3, YY3 = np.meshgrid(x, y, indexing='ij', sparse=True)
print("ij+sparse:", XX3.shape, YY3.shape)   # (5, 1) (1, 3)

出力から3つのことが読み取れます。第一に、既定の indexing='xy' では 形状が $(n_y, n_x)$ と反転 します。これは画像やプロットの慣習(行=縦=y、列=横=x)に合わせるためです。第二に indexing='ij' なら $(n_x, n_y)$ と、引数の順どおりの形状になります。行列の添字 $(i,j)$ の感覚に合うのはこちらです。第三に、sparse=True を付けると $(5,1)$ と $(1,3)$、つまりまさにブロードキャスト用の列ベクトルと行ベクトルが返ってきます。

ここが本質です。np.meshgrid(..., sparse=True) は、x[:, None]y[None, :] を作る関数にほかなりません。

import numpy as np

x = np.linspace(-2, 2, 5)
y = np.linspace(-1, 1, 3)

# meshgrid(sparse=True, indexing='ij') と手書きの None は完全に等価
XX, YY = np.meshgrid(x, y, indexing='ij', sparse=True)
print(np.array_equal(XX, x[:, None]), np.array_equal(YY, y[None, :]))

# 格子上で f(x,y) = x^2 + y^2 を評価する
F1 = x[:, None] ** 2 + y[None, :] ** 2       # ブロードキャストで (5,3)
F2 = XX ** 2 + YY ** 2
print(F1.shape, np.allclose(F1, F2))

両方の判定が True になり、形状は $(5,3)$ で一致します。indexingsparse の効果を1枚にまとめておきます。

左と中央にindexing='xy'と'ij'の格子点を同じ位置に描き赤い番号でメモリ上の並び順の違いを示し、右にsparse=Trueが返す(4,1)と(1,3)の配列を示した図

左と中央で 点の位置はまったく同じ なのに、赤い番号と線がたどる順路がまるで違います。'xy' は x方向に進んでから行が変わり、'ij' は y方向に進んでから列が変わる。形状も $(3,4)$ と $(4,3)$ で入れ替わります。右のパネルが示すとおり、sparse=True が返すのは結局 x[:, None]y[None, :] の2本だけで、密な格子を2枚作る必要はどこにもありません。

つまり 等高線を描くために密な座標行列を作る必要はまったくない のです。np.meshgrid(x, y) は $n_x n_y$ 個の要素を持つ配列を2つ実体化しますが、sparse=TrueNone を使えば $n_x + n_y$ 個で済みます。1000×1000の格子なら、8MB×2 が 8KB×2 になる計算です。

使い分けの指針

目的 推奨 理由
plt.contour(X, Y, Z) に渡す meshgridxy、密) matplotlibの多くのAPIが密な座標行列を期待する
関数値を計算するだけ x[:, None], y[None, :] メモリが $O(n_x + n_y)$ で済む
3次元以上の格子 np.ix_ または None の重ね meshgrid は次元が増えると読みにくい
インデックス的な格子 np.ogrid / np.mgrid スライス記法で簡潔に書ける

なお np.meshgridindexing は、間違えると 転置した図が出る という形で表面化します。plt.contourf(X, Y, Z)indexing='ij' の配列を渡すと、x軸とy軸が入れ替わった絵になります。「等高線が90度回っている」ときは、まずここを疑ってください。

外積とufuncのouter

$(3,1) \times (1,4)$ が外積になることは既に見ました。NumPyには専用関数もあります。

import numpy as np

u = np.arange(1, 4)
v = np.arange(1, 5)

print(np.array_equal(u[:, None] * v[None, :], np.outer(u, v)))    # True
print(np.array_equal(u[:, None] + v[None, :], np.add.outer(u, v)))  # True
print(np.add.outer(u, v).shape)   # (3, 4)

すべて True および (3, 4) が出力されます。np.outer は掛け算専用ですが、すべてのufuncには .outer メソッドがありnp.add.outernp.subtract.outernp.maximum.outer のように全組み合わせ演算を書けます。とはいえ、u[:, None] op v[None, :] のほうが「何をしているか」が形状から明白なので、私は None 記法を推します。3項以上に拡張するときも a[:,None,None] + b[None,:,None] + c[None,None,:] と自然に伸びます。

3パターンを押さえたところで、ブロードキャストの最大の落とし穴 — メモリの話に進みます。

メモリの落とし穴: コピーしないのに巨大になる理由

もう一度確認します。ブロードキャストは 入力をコピーしません(stride 0)。しかし 出力は必ず実体化されます。この非対称性が、実務でもっとも痛い事故を生みます。

距離行列の例で見積もってみましょう。$N$ 個と $M$ 個の $D$ 次元点について、X[:, None, :] - Y[None, :, :] が作る中間配列のバイト数は

$$ \text{中間配列のバイト数} = N \times M \times D \times 8 \ \ [\text{byte}] \quad (\text{float64}) $$

です。さらに結果の距離行列そのものが $N \times M \times 8$ バイト必要です。具体的な数字を入れてみます。

$N$ $M$ $D$ 中間 $(N,M,D)$ 結果 $(N,M)$
1,000 1,000 3 24 MB 8 MB
3,000 3,000 3 216 MB 72 MB
10,000 10,000 3 2.4 GB 800 MB
20,000 20,000 3 9.6 GB 3.2 GB
5,000 5,000 128 25.6 GB 200 MB

この表を、$N$ と $D$ の両方向へ広げてプロットすると次のようになります。

左に点数Nに対する必要メモリの両対数プロット(D=3,32,128の中間配列と展開式版の(N,N)、8GBの目安線つき)、右に8GBに収まる最大のNをDごとに比較した棒グラフを並べた図

左の両対数プロットでは、どの線も傾き2の直線($N^2$ に比例)になり、$D$ を大きくすると青・紫・赤の順に上へ平行移動します。展開式版(緑の破線)だけが $D$ に依存しません。右の棒グラフはこれを「8GBに収まる最大の $N$」に翻訳したもので、ブロードキャスト1行では $D=3$ で約18,000点、$D=128$ で約2,800点、$D=512$ では約1,400点しか扱えません。一方で展開式はどの $D$ でも約31,600点で一定 です。次元が上がるほど、2つの書き方の差が開いていくことがはっきり読み取れます。

表の最下段が典型的な事故です。128次元の埋め込みベクトル5000本の距離行列は、結果はたった200MBなのに、中間配列が25.6GBを要求します。ノートPCなら確実にスワップかクラッシュです。ところが展開式を使えば、中間の最大サイズは $(N,M)$ の200MBで済みます。数学的に等価な2つの式が、メモリ消費で100倍以上違うわけです。

実測してみましょう。

import numpy as np, tracemalloc, time

rng = np.random.default_rng(0)
N = M = 3000; D = 3
X = rng.normal(size=(N, D)); Y = rng.normal(size=(M, D))

tracemalloc.start()
t = time.perf_counter()
Db = np.sqrt(((X[:, None, :] - Y[None, :, :]) ** 2).sum(-1))
tb = time.perf_counter() - t
peak_b = tracemalloc.get_traced_memory()[1]; tracemalloc.stop()
print("ブロードキャスト版: ピーク %.0f MB, %.3f s" % (peak_b / 1e6, tb))
del Db

tracemalloc.start()
t = time.perf_counter()
s = (X**2).sum(1); t_ = (Y**2).sum(1)
Dg = np.sqrt(np.maximum(s[:, None] + t_[None, :] - 2*(X @ Y.T), 0))
tg = time.perf_counter() - t
peak_g = tracemalloc.get_traced_memory()[1]; tracemalloc.stop()
print("展開式版          : ピーク %.0f MB, %.3f s" % (peak_g / 1e6, tg))
print("理論値: 中間 %.0f MB, 結果 %.0f MB" % (N*M*D*8/1e6, N*M*8/1e6))

手元の環境では、ブロードキャスト版がピーク288MB・0.14〜0.22秒、展開式版がピーク144MB・0.03〜0.05秒でした(時間は実行のたびに多少ぶれますが、ピークメモリは何度実行しても288MBと144MBで完全に一定です)。理論値は中間216MB・結果72MBなので、ブロードキャスト版の288MB $= 216 + 72$ はぴったり「中間配列+結果」に一致します。展開式版の144MB $= 72 \times 2$ は「行列積の結果+距離行列」ぶんです。理論的な見積もりが実測とここまで正確に合う ことは、逆に言えば「事前に計算すれば事故は防げる」ことを意味します。大きな配列を扱う前に $N \times M \times D \times 8$ を暗算する習慣をつけましょう。

チャンク化という保険

展開式が使えない演算(一般のカーネル、Lpノルム、複雑な相互作用など)では、行を分割して処理する チャンク化が有効です。

import numpy as np

def pairwise_dist_chunked(X, Y, chunk=512):
    """行方向にchunk行ずつ処理してピークメモリを抑える"""
    N, M = X.shape[0], Y.shape[0]
    out = np.empty((N, M), dtype=float)
    for s in range(0, N, chunk):
        e = min(s + chunk, N)
        d = X[s:e, None, :] - Y[None, :, :]      # (chunk, M, D)
        out[s:e] = np.sqrt((d ** 2).sum(-1))
    return out

rng = np.random.default_rng(0)
X = rng.normal(size=(2000, 4)); Y = rng.normal(size=(1500, 4))
ref = np.sqrt(((X[:, None, :] - Y[None, :, :]) ** 2).sum(-1))
print("最大誤差:", np.abs(pairwise_dist_chunked(X, Y) - ref).max())

最大誤差は 0.0 になります。分割しても結果は完全に同一で、ピークメモリだけが $N \times M \times D \times 8$ から $\text{chunk} \times M \times D \times 8$ へ、この例では4分の1以下に下がります。チャンクサイズは、L2/L3キャッシュに収まる程度(数百行)にすると速度面でも有利になることが多いです。「1行で書ける」ことと「1行で書くべき」ことは別問題 — ブロードキャストを使いこなすとは、この判断ができるようになることでもあります。

最後に、実際に手を動かすとき必ず出会うエラーメッセージの読み方を整理しておきます。

エラーメッセージの読み方

ブロードキャストのエラーは、慣れれば5秒で原因が特定できます。パターンは3つしかありません。

タイプA: operands could not be broadcast together

ValueError: operands could not be broadcast together with shapes (3,4) (3,2)

もっとも一般的な形です。読み方は単純で、括弧内の2つの形状を右揃えして、上下で見比べる だけです。

(3, 4)
(3, 2)
       ^ ここが 4 vs 2。一致もせず、どちらも1でもない → 規則3違反

対処は「本当に揃えたかった軸はどれか」を考え直すことです。$(3,4)$ と $(3,)$ のエラーなら、右揃えの結果 $(1,3)$ と解釈されているので、b[:, None] として $(3,1)$ にするのが正解、というように進みます。

タイプB: non-broadcastable output operand

ValueError: non-broadcastable output operand with shape (4,) doesn't match
            the broadcast shape (3,4)

これは in-place演算+=*= など)に特有のエラーです。

import numpy as np

A = np.zeros((3, 4)); b = np.zeros((4,))
A += b        # OK: 出力先 A の形状 (3,4) は結果の形状と一致
try:
    b += A    # NG: 出力先 b は (4,) だが結果は (3,4)
except ValueError as e:
    print("エラー:", e)

A += b は通り、b += A は上記のエラーになります。in-place演算では 書き込み先の形状を変えられない ので、ブロードキャスト後の形状が書き込み先と一致していなければなりません。ブロードキャスト自体は成立するのに、代入先が小さすぎる — というのがこのエラーの意味です。対処は b = b + A と新しい配列を作るか、そもそも設計を見直すかです。

このエラーは、「a と b の順序を入れ替えたら通った」という形で気づきにくいバグにつながる ことがあります。+= を使うときは常に「左辺の形状は結果の形状と同じか?」を確認しましょう。

タイプC: shape mismatch: objects cannot be broadcast to a single shape

ValueError: shape mismatch: objects cannot be broadcast to a single shape.
            Mismatch is between arg 0 with shape (3, 4) and arg 1 with shape (2,).

np.broadcast_shapesnp.broadcast_arrays、あるいは3つ以上の配列を扱う関数で出ます。親切なことに どの引数同士が衝突したか(arg 0 と arg 1) まで教えてくれます。多引数の場合は、この情報を頼りに絞り込みます。

デバッグの手順

エラーが出たときの手順を定型化しておきましょう。

  1. 形状を全部printする。 print(A.shape, b.shape) を挟むのが最短経路です。頭の中で追うより確実です。
  2. 右揃えして紙に書く。 短いほうの左に1を補い、軸ごとに縦に見比べます。
  3. 「本当はどの軸を対応させたかったか」を言語化する。 「$N$ 個のサンプルそれぞれに、$D$ 個の特徴量ごとの平均を引きたい」のように日本語で書くと、必要な形状が自ずと決まります。
  4. np.broadcast_shapes で事前確認する。 実配列を作らず形状だけ検証できるので、巨大配列の前に必ず通します。
import numpy as np

print(np.broadcast_shapes((5, 1, 3), (1, 4, 1), (3,)))   # (5, 4, 3)
try:
    np.broadcast_shapes((3, 4), (2,))
except ValueError as e:
    print("事前検出:", e)

1行目は3つの配列のブロードキャスト結果 $(5,4,3)$ を、実際の配列を1つも作らずに返します。2行目は巨大なメモリを確保する前に失敗を検出できています。形状の設計は、計算を走らせる前に済ませる — これがブロードキャストと付き合う最良の作法です。

もう1つ、エラーにならないが危険なケースにも触れておきます。前のセクションで見た「正方行列を $(n,)$ で割ると静かに間違う」問題です。防御策は2つあります。1つは keepdims=True を徹底すること。もう1つは、重要な箇所で assert Z.shape == (N, D) のように形状をアサートすることです。特に他人が使う関数を書くときは、入口で形状を検査しておくと事故が激減します。

ここまでで規則・実装・落とし穴が揃いました。最後に、これらをまとめて可視化と速度測定で確かめます。

ベンチマーク: ブロードキャストは本当に速いのか

規則も内部実装も押さえたので、最後に「結局どれだけ速いのか」を実測します。比較するのは4通りの書き方です。素朴な二重ループ、外側だけループを残す半ベクトル化、フルブロードキャスト、そして展開式(行列積)。題材は800点と600点の3次元データ、つまり48万ペアの距離行列です。

計測でよくある失敗は、1回しか測らずに結論を出してしまうことです。ミリ秒オーダーの処理はOSのスケジューリングやキャッシュの状態で簡単に数倍ぶれますし、BLASを使う展開式は初回呼び出しにスレッドプールの初期化が乗ります。そこで、ウォームアップを1回入れたうえで、速い3つは5回まわして最速値を採ることにします。

import numpy as np, time

rng = np.random.default_rng(0)
N, M, D = 800, 600, 3
X = rng.normal(size=(N, D)); Y = rng.normal(size=(M, D))
_ = np.random.rand(50, 3) @ np.random.rand(3, 50)     # BLAS のウォームアップ

def best_of(fn, rep=5):
    """rep 回まわして最速値を採る(実行ごとのぶれを抑えるため)"""
    out, best = None, float("inf")
    for _ in range(rep):
        t0 = time.perf_counter()
        out = fn()
        best = min(best, time.perf_counter() - t0)
    return out, best

def loop_full():                       # 素朴な二重ループ
    R = np.empty((N, M))
    for i in range(N):
        for j in range(M):
            d = X[i] - Y[j]
            R[i, j] = np.sqrt(d @ d)
    return R

def loop_half():                       # 内側だけベクトル化(外側は Python ループ)
    R = np.empty((N, M))
    for i in range(N):
        R[i] = np.sqrt(((X[i] - Y) ** 2).sum(1))
    return R

s = (X ** 2).sum(1); u = (Y ** 2).sum(1)
D1, t_loop = best_of(loop_full, rep=1)
D2, t_half = best_of(loop_half)
D3, t_bc   = best_of(lambda: np.sqrt(((X[:, None, :] - Y[None, :, :]) ** 2).sum(-1)))
D4, t_gemm = best_of(lambda: np.sqrt(np.maximum(s[:, None] + u[None, :] - 2 * (X @ Y.T), 0)))

for name, t in [("二重ループ", t_loop), ("半ベクトル化", t_half),
                ("ブロードキャスト", t_bc), ("展開式(行列積)", t_gemm)]:
    print("%-16s : %.4f 秒   二重ループの %6.1f 倍速" % (name, t, t_loop / t))
print("一致確認 最大絶対誤差:", np.abs(D1 - D3).max(), np.abs(D1 - D4).max())

手元の環境での測定例は、二重ループが約0.43秒、半ベクトル化が約0.0054秒、ブロードキャストが約0.0060秒、展開式が約0.0009秒でした。倍率にすると順に1倍・約81倍・約72倍・約485倍です。絶対値はマシンによって変わりますが、オーダーの関係(二重ループだけが2桁遅い)は環境を選ばず再現します。一致確認の最大絶対誤差はどれも $10^{-14}$ 以下で、4つの実装が同じ行列を返していることも確かめられました。

左に4通りの実装の実行時間を対数軸の棒グラフで比較(二重ループ0.4334秒に対し他は0.001〜0.006秒)、右に二重ループを1倍としたときの速度倍率を示した実測ベンチマーク図

対数軸で見ると、二重ループだけが飛び抜けて高い柱になっているのが一目で分かります。ここから読み取れることは3つあります。

第一に、二重ループとの差は2桁です。Pythonのループ1周あたりに走るオブジェクト生成・型チェック・関数呼び出しのオーバーヘッドが、48万回積み上がった結果がこの差です。

第二に、半ベクトル化(外側だけループ)でもフルブロードキャストとほぼ同じ速度が出ています。むしろこの測定では半ベクトル化のほうがわずかに速く出ました。ボトルネックは「Pythonのループ回数」であり、$N \times M = 48$ 万回を $N = 800$ 回まで減らした時点で問題の大半は解決しているのです。「ベクトル化=常に最速」ではない — フルブロードキャスト版は $(N,M,D)$ の中間配列を確保・書き込み・読み出しするコストを払っており、ループ削減の利得をメモリ帯域が食っています。

第三に、展開式(行列積)が頭ひとつ抜けて速いという点です。$\bm{X}\bm{Y}^\top$ は最適化されたBLASのGEMMに落ちるため、キャッシュブロッキングやSIMDが効きます。中間配列を作らないぶんメモリ帯域も食いません。速度でもメモリでも有利なので、$D$ が大きいときは第一候補になります。

ただし、この4本のうち 展開式だけは測定値のぶれが桁で効きます。$D=3$ では $\bm{X}\bm{Y}^\top$ 自体が数十マイクロ秒しかかからず、実行時間の大半は $(N,M)$ 配列の確保とBLASのスレッド起動が占めるからです。実際、同じコードを繰り返すと展開式は0.0009秒から0.005秒程度まで揺れ、負荷の高いマシンではブロードキャスト版より遅く出ることさえあります(二重ループの0.43〜0.53秒と、残り3本が0.01秒未満という桁の関係は毎回再現します)。一方 $D=64$ で同じ測定をすると、ブロードキャスト版が約0.025秒に悪化するのに対し展開式は約0.003秒のままで、約10倍の差が安定して出ます。展開式の優位は「$D$ が大きいほど、そして中間配列が効いてくるほど、はっきりする」と理解するのが正確です。

まとめると、最適化の優先順位は「①Pythonループを消す → ②中間配列を作らない書き方を探す → ③どうしても必要ならチャンク化する」の順です。①だけで2桁、②でさらに数倍、というのが典型的な効き方になります。

まとめ

本記事では、NumPyのブロードキャストを規則・実装・実用パターン・落とし穴の4方向から解説しました。

  • 3つの規則がすべて — ①形状を右揃えして足りない側の左に1を補う、②軸ごとに「一致」か「片方が1」なら通し、1のほうを伸長する、③それ以外はエラー。長さ1でない軸を繰り返すことは決してありません。
  • 長さ1の軸は「インデックスを無視する軸」 — $\iota(i;1)=0$ という写像がすべてを説明します。$(3,1)+(1,4)$ が $C[i,j]=a_i+b_j$ になるのは、$\bm{a}$ が $j$ を無視し、$\bm{b}$ が $i$ を無視するからです。
  • 自動補完は左だけ、右は手動x[None, :] は省略できますが、x[:, None] は自分で書かないと絶対に起きません。行方向に配りたいときは [:, None]keepdims=True が必須です。
  • 伸長は stride 0 で実現される — 入力はコピーされず、同じメモリを何度も読むだけです。だから速く、キャッシュ効率もよい。ただし nbytes は実消費を反映せず、broadcast_to の結果は書き込み禁止です。
  • 出力は必ず実体化される — ここが最大の落とし穴。$(N,M,D)$ の中間配列は $N M D \times 8$ バイトを要求します。$\|\bm{x}-\bm{y}\|^2 = \|\bm{x}\|^2 – 2\bm{x}^\top\bm{y} + \|\bm{y}\|^2$ の展開式や、チャンク化で回避しましょう。
  • 速度の本質はPythonループの削減 — 二重ループとの差は2桁ですが、フルブロードキャストが半ベクトル化より速いとは限りません。メモリ帯域とのトレードオフを意識してください。
  • エラーは3種類operands could not be broadcast together(形状不一致)、non-broadcastable output operand(in-placeの書き込み先が小さい)、shape mismatch(多引数)。まず形状をprintし、右揃えして紙に書くのが最短経路です。

ブロードキャストを身につけると、NumPyのコードは驚くほど短くなります。しかし短さは目的ではありません。「この形状の配列とこの形状の配列を組み合わせると、意味的に何が起きるのか」を形状だけから読み取れるようになること が、本当のゴールです。そこに到達すれば、他人の書いた X[:, None, :] - Y[None, :, :] を見た瞬間に「ああ、全ペアの差ベクトルを作っているな」と即座に理解できるようになります。

次のステップとして、以下の記事も参考にしてください。