[Lucid] 텐서 Operation 시스템 구축

안암동컴맹·2025년 11월 22일
Lucid

Lucid Development

목록 보기
3/20
post-thumbnail

⚙️ 텐서 Operation 시스템 구축

🌱 Tensor 클래스의 구조적 틈새 매꾸기

Lucid의 Tensor가 어느 정도 형태를 갖추기 시작했을 때, 나는 한 가지 불편한 점이 눈에 띄었다. Tensor 자체는 데이터를 품고, gradient를 추적하고, 그 역전파를 수행할 수 있을 만큼 튼실해졌지만, 막상 연산 자체가 이루어지는 순간을 떠올리면 내부의 모든 구성물이 일관적이고 우아하게 움직이지 않는다는 사실 이 너무 선명하게 느껴졌다.

add, sub, mul, matmul 처럼 가장 기본적인 연산들을 시도하는 순간, Tensor 클래스가 아무리 탄탄하더라도 연산 과정 자체에서는 반복되는 패턴과 일종의 보일러플레이트(boilerplate) 코드가 노골적으로 드러났다.

  • 결과값 Tensor 하나를 생성하고,
  • 부모 노드를 기록하고,
  • backward 시 어떤 gradient를 부모들에게 흘려보낼지 정의하고,
  • 만약 broadcasting이 되어있으면 그 shape을 줄여 맞추고,
  • requires_grad 여부에 따라 backward 연결 여부를 결정하고...

이런 모든 과정은 Tensor가 아니라 바로 "연산(operation)" 이라는 존재가 스스로 해야 하는 일처럼 보였지만, 정작 그 책임을 어디에도 둔 적이 없었다.

이번 기록은 그런 공백을 채우기 위한 Lucid의 세 번째 핵심 축인, 텐서 연산 추상화를 위한 operation 클래스와 이를 감싸는 @func_op 데코레이터 팩토리의 형성 과정을 담고 있다. 앞서 설명한 배경 속에서 어떻게 Lucid가 연산을 바라보는 방식을 형식화하고, autodiff 그래프를 일관적으로 구축하는 구조를 만들었는지를, 가능한 한 내적 고민과 함께 서술하고자 한다.

⚠️ 초창기 텐서 연산 구현에서의 비효율성

Lucid의 가상 초창기 add 연산은 지금 다시 보면 다소 투박하고 원시적이다.

def add(a: Tensor, b: Tensor) -> tuple[Tensor, Callable]:
	ret = Tensor(a.data + b.data)
    ret.requires_grad = a.requires_grad or b.requires_grad
    
    def backward(...) -> tuple[NumPyArray, NumPyArray]:
        # d(ret)/da, d(ret)/db
    	return ret.grad, ret.grad
    
    ret._prev = [a, b]
    ret._op_backward = backward  # 함수 바인딩
    
    return ret, backward

하지만 당시에는 충분히 합리적인 출발점이었다. 무엇보다 중요한 것은, Tensor 내부에 수작업으로 메타데이터를 붙이고 backward 규칙을 그 자리에서 정의해버리는 방식이 처음에는 꽤나 단순하게 느껴졌다는 점이다.

그러나 연산이 늘어날수록, 이 단순함은 점점 무거운 의무감으로 변해갔다.

연산을 하나 만들 때 마다 나는 언제나 같은 문장을 반복해야 했다. requires_grad를 전파하고, 결과 Tensor를 만들고, 그래프 연결을 하고, backward를 등록하고... 이 모든 과정이 마치 일종의 일상적인 의식처럼 반복되었다.

Tensor는 이미 충분히 제 역할을 나름대로 잘 하고 있는데, 그 텐서들이 실제 연산을 통해 변형될 때의 패턴은 전혀 추상화되지 않은 채로 방치되어 있었다. 그래프는 분명 존재하는데, 그 그래프의 노드를 어떻게 만들지에 대한 일관된 규약이 없던 것이나 마찬가지이다.

이러한 상태가 지속되면 코드를 유지 보수하는 과정에서 실수나 누락이 생기기 마련이고, 이는 딥러닝 프레임워크로서의 Lucid의 핵심을 뒤흔드는 잠재적인 취약점이 되었다.

이 반복을 끊기 위해선 반드시 연산을 연산답게 체계화하는 장치가 필요했다.

🔮 Operation 객체라는 사고의 전환

문제의 근본을 파헤치기 위해, 나는 Tensor 중심 사고에서 벗어나 연산 자체를 객체로 만들자는 결론에 도달했다. 정확하게 말하자면, operation이라는 부모 클래스 를 체계화 한 뒤, 각 연산들(add, sub 등)을 이 operation을 상속하는 자식 클래스 로써 구현하여 텐서 연산을 추상화하는 아이디어를 떠올렸다.

연산(operation)을 하나의 작은 작업 단위라고 생각해보면 그 작업이 하는 일은 크게 두 가지 뿐이다.

1. Forward 계산을 수행한다.
2. Backward 계산 시 필요한 정보를 저장하고, gradient를 역으로 계산한다.

이 두 축만 명확히 분리된다면, 나머지 Tensor에 관련된 메타데이터 작업은 별도의 루틴을 만들면 될 것이라고 생각했다. 그래서 처음에는 다음과 같은 최소한의 operation 클래스를 스케치했다.

from abc import ABC, abstractmethod

class operation(ABC):
	def __init__(self, *args, **kwargs) -> None: ...
    
    @abstractmethod
    def forward(self, *args) -> Tensor | tuple[Tensor, ...]: ...
    
    @abstractmethod
    def backward(self, grad_out) -> Tensor | tuple[Tensor, ...]: ...
    
    def __call__(self, *args) -> Tensor | tuple[Tensor, ...]:
    	return self.forward(*args)

이 프로토타입은 나름 훌륭했다. 연산을 forward/backward 두 관점에서 구조화할 수 있었고, 각 연산마다 필요로 하는 별도의 속성(attribute)들을 따로 저장할 수도 있었다. 하지만 곧 또 다른 중복이 드러났다. 여전히 결과 Tensor를 만드는 과정, _prev, _op를 세팅하는 과정, backward 클로져(closure)를 정의하는 과정 등의 일련의 공통된 로직을 모든 연산 클래스 구현자가 직접 작성해야 한다는 문제가 남아 있었다.

즉, operation 클래스라는 형식은 forward/backward를 담는 그릇 으로서 훌륭했지만, Tensor 그래프 연결이라는 본질적 반복 문제는 여전히 해결되지 않았다.

operation은 최신 버전의 Lucid에서는 많이 다른 형태를 띄고 있으며, lucid._backend.core.operation에 존재하며 .core namespace는 생략 가능하다.

🔬 @func_op 데코레이터 팩토리의 등장

결국 나는 "연산이 해야 하는 것은 forward/backward 뿐이다" 라는 명제를 다시 떠올렸다. 그 외의 모든 작업, 즉 새 Tensor에 대한 메타데이터 연결, backward closure 구성, shape 정렬 같은 수작업들은 연산자에게 떠넘길 필요가 전혀 없었다.

그 순간 문득 "데코레이터(decorator)" 라는 개념이 떠올랐다.

연산자는 오직 실제로 계산되는 forward 함수와 그에 대응하는 gradient 계산 함수만 정의하고, 앞서 언급한 나머지 모든 Tensor 관련 작업들은 데코레이터가 자동으로 수행 한다면 어떨까?

이 발상은 곧바로 @func_op라는 데코레이터 팩토리(decorator factory) 의 탄생으로 이어졌다. 데코레이터는 연산 함수의 책임을 좁히고, 그외의 요소는 모두 공통 템플릿으로 끌어안는다. 마치 연산 구현자가 수학적 정의만 써놓으면 Lucid가 그래프 연결을 대신 다 처리해주는 형태가 되는 것이다.

이때 나는 연산 정의자가 반드시 다음 형식을 return 하도록 규약을 세웠다.

((result_tensor, grad_function), ...)

여기서 grad_function은 backward 시 호출되며, 입력 텐서 개수에 맞춘 gradient 튜플을 돌려준다. 예를 들어 어떤 operation ∗\ast가 x,yx,y를 입력으로 받아 zz를 return 한다면 (x∗y=zx\ast y=z), 이 operation은 하나의 (result_tensor, grad_function) 튜플을 return하고, 이 grad_function은 ∂z/∂x\partial z/\partial x와 ∂z/∂y\partial z/\partial y 튜플을 return한다. 이 형태만 통일해두면 나머지는 데코레이터가 어떻게는 처리할 수 있다.

@func_op은 lucid._backend.core.func_op에 존재하며, .core namespace는 생략해도 마찬가지로 접근이 가능하다.

🧬 @func_op 내부 알고리즘 파헤치기

이제 실제로 @func_op이 실제로 어떻게 생겼는지를 큰 덩어리 하나로 보는 대신, 작은 code snippet들로 잘게 분해해서 살펴보려 한다. 각 snippet은 당시 내가 어떤 고민을 했는지, 그리고 이 추상화가 어떤 문제를 해결하려고 했는지를 그대로 반영하고 있다.

가장 바깥쪽 골격은 다음과 같다.

import functools

def func_op(n_in: int, n_ret: int | None, has_gradient: bool = True) -> Callable:
	def decorator(func: Callable) -> Callable:
    	@functools.wraps(func)
        def wrapper(op_self: operation, *args, **kwargs) -> Tensor | tuple[Tensor, ...]:
        	...
        return wrapper
    return decorator

겉으로 보면 단순한 decorator factory이다. 하지만 이 안쪽 wrapper가 Lucid의 모든 operation 호출을 관통하는 공통 템플릿이 된다.

  • n_in은 연산의 인자들 중 앞에서부터 몇 개를 Tensor로 취급할 것인지를,
  • n_ret은 이 연산이 몇 개의 Tensor를 return하는지를 나타낸다.
  • has_gradient는 말 그대로 이 연산이 미분 가능한 연산인지 여부를 나타내는 플래그이다.

이 세 개의 값만으로 wrapper는 연산의 형태 를 이해하고, 나머지 세부 구현은 func라는 callback에 위임하는 방식이다.

1️⃣ wrapper 내부 상태 초기화

wrapper가 호출되면 가장 먼저 Tensor 인자들을 모아둘 컨테이너와 requires_grad 플래그를 초기화한다.

def wrapper(op_self: operation, *args, **kwargs) -> Tensor | tuple[Tensor, ...]:
	tensors: tuple[Tensor] = ()
    requires_grad = False
    ...

여기서 tensors는 나중에 backward 때 gradient를 흘려보내야 할 부모 Tensor들을 순서대로 담는 튜플이다. 이 순서가 중요한 이유는, grad function이 반환하는 gradient 튜플 역시 입력 Tensor와 동일한 순서를 따라야 하기 때문이다.

requires_grad는 모든 입력 Tensor를 훑은 뒤 "이 연산 결과가 gradient를 계산해야 하는지" 를 결정하는 전역적인 플래그가 된다. 이 두 변수는 wrapper의 나머지 로직이 모두 공유하는 작은 state들이다.

초기 구현에서는 이 두 변수를 따로 두지 않고, 인자 처리와 gradient 전파를 섞어서 처리하려고 했었다. 하지만 곧 읽기도 어렵고, 디버깅을 하기도 힘들다는 사실을 깨닫고 이처럼 wrapper의 시작 부분에서 공통 state를 분리해 선언하는 방식으로 정리했다.

2️⃣ Tensor 인자 구간 잘라내기

다음 snippet은 어디까지를 Tensor 인자로 볼 것인가를 결정하는 부분이다.

def wrapper(op_self: operation, *args, **kwargs) -> Tensor | tuple[Tensor, ...]:
	...
    if n_in is None:
    	tensor_args = args
    else:
    	if len(args) < n_in:
        	raise ValueError(
            	f"Expected at least {n_in} tensor arguments, got {len(args)}"
            )
        tensor_args = args[:n_in]
	...

n_in이 None이면 이 연산은 "앞에 오는 모든 인자" 를 Tensor로 취급한다. 예를 들어 sum(*xs, axis=None)같은 polyadic 연산에서 유용하다. 반대로 n_in이 구체적인 정수라면, 그만큼만 앞에서 잘라 tensor_args로 간주하고 나머지는 axis, keepdims, shape 처럼 연산의 모양을 제어하는 부가적인 키워드 인자로 본다.

중요한 점은, 이 시점에서 아직 tensor_args의 원소들이 실제로 Tensor인지 아닌지는 확인하지 않는다는 것이다. 이 단계의 목적은 오직 "위치 기반으로 Tensor 후보들을 분리해 내는 것" 이다.

타입 강제 casting은 다음 단계에서 일괄적으로 수행된다. 이렇게 구간 분리와 타입 정제를 나눠둔 이유는, 나중에 함수 signature을 바꾸거나 인자 수를 조정할 때 어느 부분을 수정해야 할지가 명확해지기 때문이다.

3️⃣ Tensor 강제 캐스팅과 requires_grad 전파

이제 Tensor 후보들을 실제 Tensor로 정제하는 과정이 이어진다.

def wrapper(op_self: operation, *args, **kwargs) -> Tensor | tuple[Tensor, ...]:
	...
    for arg in tensor_args:
    	tensor = lucid._check_is_tensor(arg)
        tensors += (tensor,)
        requires_grad = requires_grad or tensor.requires_grad
    ...

여기서 lucid._check_is_tensor는 이 라이브러리 전체에서 쓰이는 작은 유틸리티로, Tensor가 아닌 입력(예: int | float | NumPyArray 등)을 적절히 Tensor로 승격시키는 역할을 하며 다음과 같이 정의되어있다.

def _check_is_tensor(any: Tensor | _ArrayOrScalar) -> Tensor:
    if not isinstance(any, Tensor):
        return Tensor(any)
    return any

이 함수를 wrapper 내부에서 호출함으로써 operation 구현자는 더 이상 이 인자가 Tensor인지 아닌지를 걱정할 필요가 없어진다. 항상 Tensor라고 가정하고 operation 코드를 작성하면 되는 것이다.

반면 requires_grad는 모든 Tensor를 순회하면서 한 번이라도 True가 발견되면 끝까지 True로 유지된다. 이 값은 나중에 결과 Tensor의 requires_grad를 설정하는 데 사용된다. 즉, 부모 중 단 하나라고 gradient를 추적하고 있다면, 이 연산의 결과 역시 반드시 computation graph에 연결되어야 한다는 뜻이다. 이 규칙은 PyTorch를 포함한 대부분의 autodiff 시스템들이 채택하는 자연스러운 전파 규칙이다.

내가 이 코드를 처음 짰을 때, requires_grad를 매번 AND/OR 중 무엇으로 조합할지가 잠깐 헷갈렸던 적이 있다. 직관적으로 생각하면 "모든 부모가 True일 때만 결과도 True" 라고 오해하기 쉽지만, 실제로는 한 부모라도 추적 대상이면 결과도 gradient 경로 위에 있어야 한다. 그래서 OR을 선택하게 되었고, 이를 정리하면서 이후 연산들이 어떻게 연결되더라도 gradient 경로가 끊기지 않는다는 확신을 얻게 되었다.

4️⃣ Tensor 인자와 비-Tensor 인자 재조합

Tensor 후보들을 모두 정제했다면, 이제 나머지 인자들과 다시 합쳐서 연산 구현자에게 넘겨줄 준비를 한다.

def wrapper(op_self: operation, *args, **kwargs) -> Tensor | tuple[Tensor, ...]:
	...
    non_tensor_args = args[n_in:] if n_in is not None else ()
    new_args = (*tensors, *non_tensor_args)
    ...

여기서 핵심은, 연산 구현자가 받는 인자 시퀀스 new_args가 다음 두 가지 성질을 만족한다는 것이다.

1. 앞쪽에는 항상 정제된 Tensor들만 온다.
2. 그 뒤에 axis, keepdims, shape 등 각 연산 특유의 메타 인자 들이 그대로 이어진다.

이렇게 해두면 예를 들어 add(self, x, y, axis=None) 같은 함수를 구현할 때, x와 y는 이미 Tensor라는 것을 가정할 수 있고, axis는 추가 처리 없이 바로 사용할 수 있다.

이 작은 레이어가 생김으로써, 연산 구현부에서는 타입 체크나 Tensor wrapping 같은 잡일(?)을 전혀 신경 쓸 필요가 없게 된다.

5️⃣ 연산 구현자 호출과 리턴값 규약

준비된 인자들을 가지고 마침내 연산 구현자 func를 호출한다.

def wrapper(op_self: operation, *args, **kwargs) -> Tensor | tuple[Tensor, ...]:
	...
    func_return_pairs = func(op_self, *new_args, **kwargs)
    ...

여기서 func는 사용자가 @func_op(...)로 감싸는 실제 연산 메서드, 예를 들어 다음과 같은 형태이다.

class add(operation):
	def __init__(self) -> None:
    	super().__init__()
    
    @func_op(n_in=2, n_ret=1)
    def forward(self, a: Tensor, b: Tensor) -> _FuncOpReturnType:
    	self.result = Tensor(a.data + b.data)
        # 결과 텐서와 backward 메서드 쌍 반환
        return self.result, self.backward
    
    def backward(self) -> _GradFuncType:
    	# d(a+b)/da, d(a+b)/db 반환 (chain rule 적용 후)
    	return self.result.grad, self.result.grad

이는 실제 초창기 Lucid의 텐서 합 연산(add)의 구현체이다. 위 코드를 통해 알 수 있듯이, 구현자는 반드시 연산의 forward 메서드에서 (Tensor, grad_function) 쌍의 튜플(들)을 return해야 한다.

이 규약 덕분에 wrapper는 연산 종류에 상관없이 항상 동일한 패턴으로 결과를 풀어낼 수 있다. 여기서 grad_function은 backward 시 호출되며, 입력 Tensor 개수만큼의 NumPy array gradient를 return하는 순수 함수이다.

6️⃣ 결과 개수 정리와 단일 리턴값 처리

연산마다 리턴하는 Tensor의 개수는 다를 수 있다. 어떤 연산은 스칼라 하나만, 어떤 연산은 세 개의 텐서를 동시에 반환할 수도 있다. 이를 일반화하기 위해 다음과 같은 로직을 사용했다.

def wrapper(op_self: operation, *args, **kwargs) -> Tensor | tuple[Tensor, ...]:
	...
	if n_ret is None:
    	if not isinstance(func_return_pairs, tuple):
        	# 연산이 (결과 텐서, 미분 메서드) 튜플을 반환하지 않을 때의 예외 처리
         	raise ValueError(...)
        num_returns = len(func_return_pairs)
    else:
    	num_returns = n_ret
    
    if num_returns == 1:
    	func_return_pairs = (func_return_pairs,)
    ...

n_ret가 명시되지 않았다면 실제 리턴값 길이에서 개수를 추론한다. 이때 리턴값이 tuple이 아니면 규약 위반 으로 보고 바로 에러를 발생시킨다. 그리고 단일 리턴값인 경우에는 (out, grad) 형태로 리턴되더라고 내부적으로는 통일성을 위해 ((out, grad),)라는 1-tuple로 정규화한다.

이렇게 해두면 이후 로직에서 단일 리턴 과 다중 리턴 을 구분할 필요 없이 동일한 루프 구조로 처리할 수 있다.

7️⃣ 결과 Tensor의 메타데이터 설정

이제 wrapper는 실제 결과 Tensor들을 하나씩 순회하면서, 각 결과에 대해 그래프 메타데이터를 채워 넣는다.

def wrapper(op_self: operation, *args, **kwargs) -> Tensor | tuple[Tensor, ...]:
	...
    results: tuple[Tensor] = ()
    for result, compute_grad in func_return_pairs:
    	result.requires_grad = requires_grad and has_gradient
        result._op = op_self  # operation 인스턴스
        results += (result,)
    	...

여기서 requires_grad는 앞서 입력 Tensor들에서 계산한 전역 플래그와, 연산 자체가 미분 가능한지에 대한 has_gradient를 함께 고려한다. 이 둘을 AND 한 값이 최종적으로 결과 Tensor에 저장된다.

_op는 이 Tensor를 생성한 연산 객체, 즉 graph 상에서 이 Tensor의 생성 원인을 가리킨다. 나중에 backward를 수행할 때, Tensor는 자신을 만든 operation의 backward 메서드를 호출해 자신을 만든 부모들에게 gradient를 전파하게 된다.

이렇게 결과를 results 튜플에 쌓아두는 이유는, 여러 개의 Tensor를 리턴하는 연산에서도 일관된 리턴 형태를 유지하기 위함이다. 마지막에 단일 리턴이면 results[0]만 꺼내고, 다중 리턴이면 튜플 그대로 돌려준다.

8️⃣ Backward 클로져 구성 – Gradient의 실제 흐름

가장 중요한 부분은 backward 시 호출될 closure 을 만드는 단계이다. 이 부분이야말로 operation abstraction의 핵심이라고 해도 과언이 아니다.

def wrapper(op_self: operation, *args, **kwargs) -> Tensor | tuple[Tensor, ...]:
	...
   	for result, compute_grad in func_return_pairs:
    	...
        def _backward_op(*, _func: Callable = compute_grad) -> None:
        	grads = _func()
            if n_in == 1 or not isinstance(grads, tuple):
            	grads = (grads,)
            
            if len(grads) != len(tensors):
            	# 리턴된 미분값의 개수가 입력 텐서 인자 수와 안맞을 때
            	raise ValueError(...)
            
            for tensor, grad in zip(tensors, grads):
            	new_grad = lucid._match_grad_shape(tensor.data, grad)
                lucid._set_tensor_grad(tensor, new_grad)
            
            return
        ...

이 함수는 나중에 Tensor의 backward 루틴에서 호출된다. 내부에서 하는 일은 다음과 같이 자연스럽게 읽힌다.

1. 먼저 연산 구현자가 정의한 compute_grad()를 호출해 원시 gradient 를 얻는다. (op_self.forward가 리턴한 미분 메서드를 지칭)

2. Unary 연산의 경우를 위해 gradient가 단일 값이면 1-tuple로 감싼다.

3. 입력 Tensor들과 gradient를 나란히 zip으로 묶어 순회한다.

4. _match_grad_shape로 broadcasting/collpase 로 인해 변한 gradient의 shape를 원래 입력 Tensor의 shape에 맞게 변화시킨다.

5. 마지막으로 _set_tensor_grad를 통해 해당 Tensor의 grad 필드에 gradient를 누적(accumulate)한다.

여기서 중요한 것은, 연산 구현자는 gradient가 수학적으로 어떤 값이어야 하는지만 정의하면 되고, 그 gradient를 어떤 순서로 어떤 Tensor에 누적해야 하는지는 decorator가 전담한다는 점이다.

참고로 lucid._set_tensor_grad 함수는 다음과 같이 구현되어 있다.

def _set_tensor_grad(tensor: Tensor, grad: _NumPyArray, at: SupportsIndex = ...) -> None:
    if not tensor.requires_grad:
        return
    if tensor.grad is None:
        tensor.grad = grad  # 미분값 최초 누적
    else:
    	# np.ndarray 타입의 tensor.grad가 read-only인 경우
        # 값 복제를 통해 강제로 writable 상태로 전환
        if not tensor.grad.flags.writeable:
            tensor.grad = tensor.grad.copy()
		
        if tensor.grad.ndim == 0:
        	# 미분값이 스칼라인 경우 단순 덧셈
            tensor.grad += grad
        else:
        	# 미분값이 다차원 배열인 경우 인덱싱 도입
            tensor.grad[at] = tensor.grad[at] + grad

그리고 gradient의 shape을 원래 Tensor의 shape과 맞춰주는 중요한 역할을 하는 lucid._match_grad_shape는 다음과 같이 구현되어 있다.

def _match_grad_shape(data: _NumPyArray, grad: _NumPyArray) -> _NumPyArray:
	# data: 원본 Tensor의 데이터 값
    # grad: 맞추고자 하는 gradient 값

    # (1) data와 grad의 shape가 이미 동일한 경우 -> 그대로 반환
    if data.shape == grad.shape:
        return grad

    # (2) data가 스칼라였던 경우
    #     forward에서 스칼라와 broadcast가 일어났다면 grad는 여러 값으로 확장되었을 것.
    #     스칼라에 대한 grad는 그 모든 요소를 더한 값이므로 sum()으로 처리.
    if data.ndim == 0:
        return grad.sum()

    # (3) grad가 스칼라인 경우
    #     forward에서 data가 grad보다 더 큰 shape였고 broadcasting 되었음.
    #     grad를 data.shape로 broadcast 시켜야 함.
    if grad.ndim == 0:
        return np.broadcast_to(grad, data.shape)

    # (4) data와 grad의 총 element 수(size)가 같다면
    #     shape만 다르고 내용물 수는 동일한 경우 -> reshape로 해결 가능
    if data.size == grad.size:
        return grad.reshape(data.shape)

    # (5) data.size > grad.size 인 경우
    #     즉, grad가 더 작은 크기 -> forward에서 grad가 반복(broadcast)되어 data가 생성됨.
    #     따라서 grad를 확장시켜야 함.
    elif data.size > grad.size:
        # grad를 flatten
        grad_squeeze = grad.flatten()
        # 얼마나 반복해야 하는지 횟수 계산
        expand_factor = data.size / grad.size

        # 반복 횟수가 정수가 아니라면 broadcast로 만들 수 없음 → 오류
        if expand_factor % 1 != 0:
            raise ValueError(
                f"Cannot broadcast grad of {grad.shape} to data of {data.shape}."
            )

        # grad를 expand_factor만큼 반복하여 확장
        grad_expand = grad_squeeze[..., None].repeat(int(expand_factor), axis=-1)

        # 원래 data의 shape로 reshape
        return grad_expand.reshape(data.shape)

    # (6) data.size < grad.size 인 경우
    #     즉, grad가 더 큰 크기 -> forward 과정에서 data가 확장(broadcast)되면서 grad가 축적되어야 함.
    #     반대로 backprop에서는 확장된 grad를 collapse 해야 함 (sum 축소)
    elif data.size < grad.size:
        # grad.size 가 data.size로 나누어 떨어지지 않으면 collapse 불가
        if grad.size % data.size != 0:
            raise ValueError(
                f"Cannot collapse grad of {grad.shape} to data of {data.shape}."
            )

        new_shape = tuple()
        remain_size = grad.size

        # data shape에 맞게 grad를 재구성할 shape을 계산
        for d_dim in data.shape:
            fac = remain_size // d_dim
            new_shape += (d_dim,)
            remain_size = fac

        # 마지막 남은 차원을 추가
        new_shape += (fac,)

        # collapse를 위해 마지막 axis를 sum하여 data.shape로 축소
        return grad.reshape(new_shape).sum(axis=-1)

    # (7) 이 외의 경우 → 처리할 수 없는 예외 상황
    else:
        raise ValueError("Unknown error occurred.")

이 유틸리티 함수의 동작을 간단하게 표로 정리하면 다음과 같다.

상황의미처리 방식
shape 동일이미 원하는 모양그대로 반환
data가 스칼라모든 grad 요소 합이 새로운 gradsum()
grad가 스칼라스칼라 →\rightarrow data.shape 로 broadcastbroadcast_to()
size 동일shape만 다름reshape()
data.size > grad.sizegrad가 반복되어 broadcastgrad를 repeat해서 확장
data.size < grad.sizegrad가 broadcast 후 축적됨grad를 reshape 후 sum()으로 collapse

생각보다 연산 과정을 거치면서 Tensor의 shape이 broadcast/collapse에 의해 변하는 경우의 수가 많아 이 함수의 로직을 구현하는 데 있어서 적지 않은 어려움을 느꼈다.

9️⃣ 결과 Tensor에 Backward 연결과 최종 리턴

마지막으로 gradient 추적 대상인 결과 Tensor에 위의 bakcward closure _backward_op를 실제로 연결하고, 리턴 값을 정리하는 단계가 온다.

def wrapper(op_self: operation, *args, **kwargs) -> Tensor | tuple[Tensor, ...]:
	...
    for result, compute_grad in func_return_pairs:
    	...
        if result.requires_grad:
        	result._prev = list(tensors)
            result._backward_op = _backward_op
        
    return results if num_returns > 1 else results[0]

result.requires_grad가 True일 때만 _prev와 _backward_op를 설정하는 이유는, gradient를 추적하지 않아도 되는 Tensor에 대해 굳이 그래프 구조를 유지할 필요는 없기 때문이다. 이렇게 함으로써 불필요한 메모리 사용을 줄이고, backward 시에도 쓸데없는 경로를 따라가지 않게 된다.

이 한 줄, result._prev = list(tensors)는 이 결과 Tensor가 그래프 상에서 어디에서 왔는지를 기록하는 핵심 연결 고리이다. 반대로 _backward_op는 이 Tensor를 기준으로 어떻게 부모들에게 gradient를 돌려줄 것인지를 encapsulate 한 함수 포인터 역할을 한다.

이 둘이 묶여 있는 구조가 바로 Lucid의 autodiff 그래프를 이루는 기본 단위이다.

그리고 마지막 return 문에서 단일 리턴과 다중 리턴을 구분해 적절한 형태로 돌려줌으로써, 사용자 입장에서는 자연스러운 Python 함수처럼 연산을 사용할 수 있게 된다.

🎁 @func_op의 특수 케이스 Wrapping

사실 대부분의 연산은 입력이 1개거나 2개였다. 다른 말로 하자면 대부분의 연산은 unary 이거나 binary 연산이라는 것이다. 그래서 이를 더 편하게 쓰기 위해 얇은 wrapper 들을 만들었다.

def unary_func_op(has_gradient: bool = True) -> Callable:
	# n_in=1, n_ret=1 고정
    return func_op(1, 1, has_gradient=has_gradient)


def binary_func_op(has_gradient: bool = True) -> Callable:
	# n_in=2, n_ret=1 고정
    return func_op(2, 1, has_gradient=has_gradient)

추가적으로 입력이 nn개인 polyadic 연산을 위한 wrapper 또한 만들었다.

def poly_func_op(has_gradient: bool = True) -> Callable:
	# n_in은 동적 추론, n_ret=1 고정
    return func_op(None, 1, has_gradient=has_gradient)

결국 이 helper들은 연산 구현자가 불필요한 고민을 하지 않게 되고, 연산 정의를 더욱 가독성 좋고 유지보수하기 쉬운 형태로 정리하기 위한 장치이다.

🪄 mul 예시를 통해 전 과정 살펴보기

이제 이렇게 구축된 Lucid의 autodiff 엔진 기반의 Tensor operation 시스템을 실제 mul 연산 예시를 통해 처음부터 빠르게 훑어보자.

0️⃣ 텐서 세팅과 원소별 곱셈 실행

a = lucid.Tensor([1, 2], requires_grad=True)
b = lucid.Tensor([3, 4], requires_grad=True)

c = a * b  # 원소별 곱셈 호출

실제로는 편의를 위해 Lucid 상의 모든 Tensor operation들은 다음과 같이 상위 namespace에서 별도의 함수로 wrapping 되어있다. 즉, 같은 덧셈 연산이어도 매 호출마다 다른 mul 클래스의 인스턴스가 생성되는 것이다. 이렇게 함으로써 그래프 상에서의 연산의 역추적이 꼬이지 않게 할 수 있다.

# lucid/_func/__init__.py
from lucid._func.bfunc import mul

def mul(a: Tensor, b: Tensor) -> Tensor:
	return mul()(a, b)  # 인스턴스 생성 후 호출(mul.__call__)

그리고 이 mul wrapping 함수는 Tensor 클래스의 기본 * 연산자 메소드인 __mul__을 override(오버라이드) 한다.

# lucid/_tensor/tensor.py
from .tensor_base import _TensorBase

class Tensor(_TensorBase): ...
# lucid/_tesnor/tensor_base.py
from typing import Self

class _TensorBase:  # Tensor 상위 클래스
	...
    # 메서드 시그니처만 정의
    def __mul__(self, other: Self) -> Self: ...
    ...
# lucid/_func/__init__.py

from lucid._tensor import Tensor

Tensor.__mul__ = mul  # 동적(런타임) 메서드 오버라이딩

이때, operation의 하위 클래스인 mul은 다음과 같이 구현되어 있다.

# lucid/_func/bfunc.py

class mul(operation):
	def __init__(self) -> None:
    	super().__init__()
    
    @binary_func_op()
    def forward(self, a: Tensor, b: Tesnsor) -> _FuncOpReturnType:
    	self.result = Tensor(a.data * b.data)
        return self.result, self.backward
    
    def backward(self) -> _GradFuncType:
    	return (
        	self.result.grad * b.data,  # d(a*b)/da = b
            self.result.grad * a.data,  # d(a*b)/db = a
        )

즉, 예시 코드에서 c = a * b를 실행하면

1. * 연산자가 a.__mul__을 접근하게 되고 이는 오버라이딩된 lucid._func.mul 함수를 실행하게 된다.

2. 그러면 lucid._func.bfunc.mul 클래스의 인스턴스를 생성하게 되고, mul()(a, b)에서 mul.__call__의 직접 호출을 통해 mul.forward(a, b)가 호출된다.

1️⃣ @binary_func_op 호출

mul.forward가 실행되기 전, 데코레이터 팩토리 wrapper @binary_func_op()에 의해 데코레이터 팩토리 @func_op()이 호출된다.

def binary_func_op(...) -> Callable:
	return func_op(n_in=2, n_ret=1, ...)

2️⃣ a, b 타입 검사 및 승격

이제 @func_op 내부 로직에 따라 연산의 입력 인자인 a, b에 대해 Tensor 인스턴스인지 타입 검사를 하고 필요에 따라 적절히 승격을 진행한다.

하지만 이 예시에서는 a, b는 이미 모두 Tensor 이므로 별도의 승격 없이 넘어간다.

3️⃣ mul.forward 호출 및 리턴 튜플 획득

이제 본격적으로 @func_op이 정제된 a, b를 mul.forward에 넘겨 실제 원소별 곱을 실행하고, 이로부터 결과 텐서와 그에 따른 backward 메서드를 리턴 받는다.

4️⃣ 결과 Tensor 후처리 및 최종 리턴

이제 @func_op이 본인의 마지막 과정인 각종 computation graph 관련 메타데이터들을 결과 Tensor에 지정한다.

  • result.requires_grad: a, b 모두 True 이므로 True
  • result._prev: 연산의 입력 인자인 a, b를 부모로 지정
  • result._backward_op: @func_op 내부에서 동적으로 생성한 closure인 _backward_op로 바인딩, 이때 closure의 _func 인자는 mul.backward 이다.

이러한 후처리가 모두 완료되면 result 텐서를 최종적으로 리턴하여 c = a * b의 실행을 미무리한다.

❇️ 결과 텐서에 대한 역전파 수행

이제 forward가 끝났으니 c.backward()를 호출하여 ∂c/∂a\partial \mathbf{c}/\partial\mathbf{a}와 ∂c/∂b\partial \mathbf{c}/\partial\mathbf{b}를 구해보자.

c.backward()

5️⃣ c.grad 초기화

c 텐서는 아무런 gradient 값을 가지고 있지 않으므로, c.backward 내부에서 최초 gradient을 1∈R2\mathbf{1}\in\mathbb{R}^2로 초기화한다.

c.grad = None -> np.ones_like(c.data)

6️⃣ 역전파 순서 설정을 위한 위상 정렬 수행

각 텐서의 backward를 호출하기 위한 순서를 정하기 위해 c._prev에서 시작하여 위상 정렬(topological sort)를 수행한다.

그에 따른 결과는 (a, b, c)가 되고, 따라서 실제 backward 호출 순서는 이의 역순 인 (c, b, a)가 된다.

이제 이 순서를 바탕으로 각 텐서의 op._backward_op이 호출되게 된다.

7️⃣ @func_op에 의해 바인딩된 c._backward_op 실행

앞서 데코레이터 @func_op에 의해 바인딩된 closure _backward_op이 실행된다.

def _backward_op(*, _func: Callable = mul.backward) -> None:
        grads = _func() # grad_a, grad_b가 담긴 튜플
        ...
        for tensor, grad in zip(tensors, grads):
            new_grad = lucid._match_grad_shape(tensor.data, grad)
            lucid._set_tensor_grad(tensor, new_grad)

mul연산은 두 개의 텐서 입력을 받으므로, _func 즉, mul.backward는 다음의 두 gradient가 담긴 튜플 (∂c/∂a, ∂c/∂b)(\partial \mathbf{c}/\partial\mathbf{a},~\partial \mathbf{c}/\partial\mathbf{b})을 리턴한다.

이후, 각 gradient의 shape이 각 input 텐서 a, b의 shape과 맞는지 확인하고 이상이 없으면 a.grad, b.grad 각각에 대응하는 gradient 값을 누적한다.

8️⃣ Non-leaf 텐서들의 Gradient 회수

실질적인 gradient 전파가 모두 완료되었으므로, 중간 단계의 non-leaf 텐서들이 가지고 있는 gradient를 다시 None으로 초기화한다.

이 예시에서는 a와 b 모두 gradient 추적 대상임과 동시에 말단 노드(terminal node) 이기 때문에 이를 제외한 c의 gradient만 회수된다(c.grad = None).

✅ 최종 결과 확인

이제 a, b, 그리고 c 텐서를 출력해보자.

Tensor([1 2], grad=[3. 4.])  # a
Tensor([3 4], grad=[1. 2.])  # b
Tensor([3 8], grad=None)  # c

이로써 간단한 forward/backward 스윙을 통해 Lucid의 핵심적인 autodiff 엔진 기반의 Tensor operation 시스템이 어떤 식으로 구동되는지 알아보았다.

🌙 Lucid의 연산 체계가 확립되다

이렇게 탄생한 @func_op 기반 operation 시스템은 단순한 기술적 장치가 아니라 Lucid 전체의 철학을 담고 있다. 연산 구현자는 수학적 정의에만 집중하고, 그래프 연결과 gradient 역전파는 자동화된다. Tensor는 데이터를 담는 그릇 이고, operation은 그릇을 변형시키는 규칙 이며, @func_op은 그 과정이 그래프 위에서 일관된 흐름으로 이어지도록 한다.

이 구조가 완성되었을 때 나는 비로소 Lucid가 autodiff 엔진으로서 하나의 유기체처럼 움직이기 시작했다는 느낌을 받았다. 연산들은 더 이상 제각각 흩어져 있는 함수가 아니라, 공통된 원리와 일관된 구조 속에서 이어진 작은 노드들의 연속이 되었다.

다음 개발 일지에서는 바로 이 abstraction 위에서 개별 연산들을 어떻게 구현했는지, 특히 add, matmul 같은 초기 연산들의 구현 과정과 그 과정에서 마주쳤던 gradient 버그들을 어떻게 하나씩 다뤄갔는지르 이어서 기록할 것이다.

profile
안암동컴맹
Korea Univ. Computer Science & Engineering
이전 포스트

[Lucid] Tensor 클래스의 구현

다음 포스트

[Lucid] 실제 텐서 연산들의 구현

0개의 댓글