[Python] Einsum

einsum 개념 및 사용법

numpy.einsum()은 아인슈타인 합(아인슈타인 표기법)을 이용하여 다차원 배열의 연산을 간결하게 수행할 수 있는 함수입니다.

1. 기본 개념: 아인슈타인 합 표기법

아인슈타인 합 표기법은 반복되는 인덱스를 축소(합산)하는 방법입니다. 예를 들어, 두 벡터 a, b의 내적(dot product)은 다음과 같이 표현됩니다.

$$ c = \sum_{i} a_{i} b_{i} $$

이 수식을 einsum으로 표현하면:

import numpy as np

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

c = np.einsum('i,i->', a, b)
print(c)

결과: \( 32 \) (벡터 내적 결과)

2. einsum의 주요 연산 예제

(1) 행렬 곱셈

행렬 곱셈(행렬 A와 B의 곱)도 einsum으로 쉽게 구현할 수 있습니다.

$$ C_{ij} = \sum_{k} A_{ik} B_{kj} $$

A = np.array([[1, 2], [3, 4]])  # (2x2)
B = np.array([[5, 6], [7, 8]])  # (2x2)

C = np.einsum('ik,kj->ij', A, B)
print(C)

출력 결과:

[[19 22]
 [43 50]]

이는 일반적인 행렬 곱셈 np.dot(A, B) 또는 A @ B와 동일합니다.

(2) 벡터 외적 (Outer Product)

외적(outer product)은 다음과 같이 einsum을 활용할 수 있습니다.

A = np.array([1, 2, 3])
B = np.array([4, 5, 6])

outer_product = np.einsum('i,j->ij', A, B)
print(outer_product)

출력 결과:

[[ 4  5  6]
 [ 8 10 12]
 [12 15 18]]

'i,j->ij'는 \( i \)와 \( j \)를 곱하여 \( (i, j) \) 형태의 2D 배열을 생성.

(3) 행렬 전치 (Transpose)

A = np.array([[1, 2, 3], [4, 5, 6]])
A_T = np.einsum('ij->ji', A)
print(A_T)

출력 결과:

[[1 4]
 [2 5]
 [3 6]]

이는 np.transpose(A) 또는 A.T와 동일합니다.

(4) 행렬 대각합 (Trace)

A = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
trace = np.einsum('ii->', A)
print(trace)  # 1 + 5 + 9 = 15

출력 결과:

15

이는 np.trace(A)와 동일한 결과입니다.

(5) 배치 행렬 곱셈 (Batch Matrix Multiplication)

3차원 이상의 행렬 곱셈도 einsum을 사용하면 쉽게 구현할 수 있습니다.

A = np.random.rand(10, 3, 4)  # 배치 크기 10, 3x4 행렬
B = np.random.rand(10, 4, 5)  # 배치 크기 10, 4x5 행렬

C = np.einsum('ijk,ikl->ijl', A, B)
print(C.shape)  # (10, 3, 5)

위 코드는 np.matmul(A, B) 또는 A @ B와 동일한 역할을 합니다.

3. einsum의 장점

• 가독성: einsum은 긴 루프 없이 수식을 직관적으로 표현할 수 있음.
• 성능 최적화: einsum은 numpy 내부에서 최적화되어 빠르게 동작하고, 불필요한 메모리 할당을 줄일 수 있음.
• 다양한 연산 지원: 내적, 외적, 전치, 대각합, 배치 연산 등 여러 수학적 연산을 하나의 함수로 수행 가능.

4. einsum 사용 팁

numpy.einsum_path()를 사용하면 연산 최적화 경로를 확인할 수 있음.

optimal_path = np.einsum_path('ik,kj->ij', A, B, optimize=True)
print(optimal_path)

최적화된 einsum 실행:

C = np.einsum('ik,kj->ij', A, B, optimize=True)

결론

• einsum은 행렬 연산을 효율적으로 수행하는 강력한 도구.
• 내적, 외적, 전치, 대각합 등 다양한 연산을 하나의 함수로 표현 가능.
• 코드 가독성을 높이고 최적화할 때 유용.

다양한 차원의 배열 연산을 직관적이고 효율적으로 구현하려면 einsum을 적극 활용하는 것이 좋습니다! 🚀