[Lucid] NumPy/MLX 통합 데이터 타입

🔮 NumPy/MLX 통합 데이터 타입

Lucid를 만들며 가장 오래 붙들고 있던 내부 모듈이 lucid.types.Numeric이었다. CPU에서는 NumPy, GPU에서는 MLX를 쓰는데, 사용자에게는 단일 dtype 경험을 제공하고 싶었다. PyTorch의 torch.dtype처럼 통합된 타입 체계를 가지면서도, MLX가 지원하지 않는 조합(예: GPU float64)을 알아서 교정하고, lazy execution 환경에서도 예측 가능한 동작을 보장하는 것이 목표였다.

이번 일지는 Lucid 고유의 데이터 타입 시스템인 Numeric을 구상한 배경, 코드 구조, GPU/CPU 매핑 로직, Tensor/Module/Random API에서의 활용, 삽질과 해결 과정을 테크 리포트 형식으로 기록한다.

---

🧭 두 개의 독립적인 dtype 시스템

Lucid는 CPU 경로로는 NumPy, GPU 경로로는 MLX를 사용한다. 두 백엔드는 dtype 이름과 제약이 미묘하게 달라 GPU에서 float64가 32비트로 내려가거나, bool 표현이 mx.bool_처럼 분리되는 일이 잦다. 사용자에게는 lucid.Float32처럼 일관된 이름을 주되 내부에서는 디바이스별 최적 dtype으로 해석되고, 파라미터 초기화부터 연산·랜덤 생성·to/from 변환까지 동일하게 적용되길 원했다.

MLX의 lazy exec과 NumPy의 즉시 실행 차이, float64 on GPU 강등 같은 이슈가 자동으로 처리되어 "사용자는 dtype 걱정을 잊고 연산에만 집중한다" 상태가 되도록 설계했다.

🧱 Numeric 클래스 구조와 속성

• 정의 위치: lucid/types.py
• 핵심 필드
  • base_dtype: int | float | complex 기본 축
  • bits: 비트 수(None이면 bit-free)
  • _np_dtype / _mlx_dtype: 실제 NumPy/MLX dtype 객체 캐시

GPU 강등 로직

MLX GPU 디폴트 장치가 float64를 완전 지원하지 않을 때, Float64 생성 시 bits를 32로 내려 _mlx_dtype을 mx.float32로 매핑한다. 이 부분이 Lucid Numeric이 단순 래퍼를 넘어 "정책"을 담는 지점이다.

class Numeric:
    def __init__(self, base_dtype, bits):
        ...
        if bits is not None:
            self._np_dtype = getattr(np, self.base_str + str(bits))
            bits_mlx = bits

            if (
                mx.default_device().type is mx.DeviceType.gpu
                and self.base_dtype is float
                and bits == 64
            ):
                bits_mlx = 32  # GPU float64 → float32 강등

            self._mlx_dtype = getattr(mx, self.base_str + str(bits_mlx))

파싱 유틸

• cpu/gpu 프로퍼티로 실제 dtype 반환.
• is_bit_free: bits가 None이면 데이터의 기존 dtype을 따라가도록 허용.
• parse(device): 디바이스별 dtype 선택.
• auto_parse(data_dtype, device): bit-free Numeric가 주어졌을 때, 실제 데이터 dtype에서 비트수를 추출해 같은 계열로 매핑.

🧲 Numeric 인스턴스 세트와 공개 API

• 프리셋 정의: lucid/types.py

  • Int, Int8, Int16, Int32, Int64
  • Float, Float16, Float32, Float64
  • Complex, Complex64

• 사용자 노출: lucid/__init__.py에서 동일한 이름으로 re-export해 lucid.Float32처럼 바로 쓰도록 했다. 타입 이름이 프레임워크 표면에 드러나야 학습 곡선이 낮아진다.

• 문자열 매핑: numeric_dict와 to_numeric_type가 dtype 문자열을 Numeric으로 변환한다. 예: np.float32 → Float32. 문자열 파서를 둔 덕분에 외부 라이브러리 dtype도 부드럽게 받아들일 수 있다.

def to_numeric_type(data_dtype: type) -> Numeric:
    str_dtype = str(data_dtype).split(".")[-1]
    name = re.findall(r"[a-z]+", str_dtype)[0]
    bits = re.findall(r"\d+", str_dtype)[0]
    return numeric_dict[name][bits]

🧬 Tensor 생성 시 dtype 흐름

• 경로: lucid/_tensor/tensor.py

• 동작 요약

  1. Python built-in 타입(int, float, complex)을 Numeric으로 매핑하기 위해 _dtype_map을 사용.
  2. dtype이 Numeric이면 Numeric.parse(device)로 디바이스별 dtype을 선택해 NumPy/MLX 배열을 만든다.
  3. 입력 데이터가 MLX array면 device를 강제로 "gpu"로 바꾸고 dtype이 맞지 않으면 astype으로 변환.
  4. device=="gpu"인데 NumPy array면 mx.array로 승격하고, bool이면 mx.bool_로 맞춘다.

Bit-Free 처리

Numeric.is_bit_free인 경우(예: Float), Tensor.astype에서 auto_parse를 호출해 데이터에 맞는 비트 폭을 추론한다(tensor.py).

def astype(self, dtype: type | Numeric) -> Self:
    new_dtype = dtype
    if isinstance(dtype, Numeric):
        self._is_bool_tensor = False
        if dtype.is_bit_free:
            new_dtype = dtype.auto_parse(self.data.dtype, device=self.device)
        else:
            new_dtype = dtype.parse(device=self.device)
    ...
    self.data = self.data.astype(new_dtype)
    self.dtype = types.to_numeric_type(self.data.dtype)

• 교훈: bit-free Numeric을 도입한 덕분에 "같은 Float인데 CPU에서는 float64, GPU에서는 float32"라는 동형 타입을 유지할 수 있었다.

🧮 Tensor.__getitem__과 dtype 보존

• 경로: tensor.py
• GPU 인덱싱 시 parse_mlx_indexing으로 bool/list 마스크를 int32로 변환한다. Numeric과 직접 연관되진 않지만, dtype 보존을 위해 GPU 경로에서도 bool→int32 변환이 일관되게 동작해야 했다.
• 슬라이스 후 새 Tensor를 만들 때 원래 dtype을 그대로 사용한다. Numeric 매핑이 깨지지 않도록 Tensor(...) 생성자에 dtype=self.dtype를 넘긴다.

🧪 Random/Functional 경로에서의 Numeric 사용

랜덤

• lucid.random.permutation 등에서 dtype: _BuiltinNumeric | Numeric를 받으며, Tensor 생성 시 Numeric으로 전달해 디바이스별 dtype을 맞춘다(lucid/random/_func.py).

Functional

• nn.functional.one_hot는 dtype: Numeric | bool | None을 받아 라벨 인코딩 결과의 dtype을 명시 가능하게 했다(lucid/nn/functional/_util.py).

🧩 Module/Buffer/Parameter와 dtype 이동

• Module.to 흐름: lucid/nn/module.py가 파라미터/버퍼/서브모듈을 재귀 이동한다. 파라미터 자체는 Numeric을 기억하고 있으므로 .to("gpu") 시 Tensor.to가 dtype을 MLX dtype으로 재해석한다.
• Buffer 등록: register_buffer(..., dtype=...) 호출 시 Numeric을 넘기면 해당 장치용 dtype으로 생성된다. MLX GPU float64 강등 규칙이 동일하게 적용된다.

🧭 PyTorch와의 호환/차이

사용 경험은 lucid.Float32, lucid.tensor(..., dtype=lucid.Float32)처럼 PyTorch의 torch.float32를 떠올리게 한다. 하지만 Lucid Numeric은 "어떤 디바이스에서 어떻게 해석될지"까지 포함해 GPU에서는 MLX 제약을 반영, float64 요청을 자동으로 float32로 매핑한다.

또 bit-free 타입(Float, Int, Complex)을 둬 데이터 주도형 bit 결정을 허용하고, Numeric.auto_parse가 lazy eval 환경에서도 올바른 비트 폭을 추출해 데이터 materialize 타이밍에 덜 민감하도록 설계했다.

🛠️ Numeric 설계 시 부딪힌 문제와 해결

MLX GPU float64 미지원

• 문제: MLX가 GPU float64를 기본 지원하지 않아, 사용자 요청이 Float64인 경우 런타임 오류가 발생.
• 해결: Numeric 생성자에서 GPU default device를 보고 bits를 32로 강등. 사용자-facing 이름은 그대로 Float64지만, 내부 GPU dtype은 mx.float32. 이 동작을 문서화하고 devlog로 명시.

Bit-Free 타입 도입 시 dtype 손실

• 문제: Float처럼 bits=None을 주면 Tensor 생성 시 어떤 비트 폭을 써야 할지 모호. 초기에는 무조건 float32로 가정했다가 CPU 사용자에게 혼란.
• 해결: Numeric.auto_parse를 만들어 입력 데이터의 dtype에서 bit를 추출. CPU에서는 float64, GPU에서는 float32가 자연스럽게 선택된다.

to_numeric_type 정규화 실패

• 문제: MLX dtype 문자열이 float32처럼 나오지 않아 regex 파싱이 실패하는 경우가 있었다.
• 해결: str(dtype)에서 모듈 경로를 잘라내고, [a-z]+ + [0-9]+ 패턴을 두 번 추출해 이름/비트를 분리. NumPy/MLX 모두 호환되는 문자열을 만들어낸다.

bool 처리

• 문제: MLX bool은 mx.bool_이고, CPU에서는 bool/np.bool_이 혼용된다.
• 해결: Tensor 생성 시 dtype이 bool이면 _is_bool_tensor 플래그를 세워 연산 경로에서 bool 전용 로직을 타게 했다(tensor.py). Numeric은 bool을 직접 다루지 않고, Tensor 레벨에서 처리한다.

🔎 Numeric이 Tensor 그래프/grad에 미친 영향

• Gradient dtype: _match_grad_shape는 data dtype과 grad dtype이 다르면 broadcast/reshape 후 덮어쓴다(lucid/__init__.py). Numeric 덕분에 CPU/GPU 모두 dtype 매칭이 쉬워졌다.

• Lazy eval: MLX에서 Tensor.eval()은 mx.eval + mx.stop_gradient를 수행(tensor.py). dtype은 이미 Numeric → MLX dtype으로 정규화되어 있어 추가 변환이 필요 없다.

• free 텐서: Tensor.free()는 디바이스 자유도를 설정하지만 dtype은 그대로 유지된다. CPU↔GPU 자동 이동 시에도 Numeric이 올바른 dtype을 선택해준다.

1️⃣ 예시: CPU→GPU 이동 시 dtype 유지

import lucid

x = lucid.tensor([1.0, 2.0, 3.0], dtype=lucid.Float64)  # CPU float64
x.device   # "cpu"
x.dtype    # lucid.types.Numeric(base_dtype=float, bits=64)

x_gpu = x.to("gpu")
x_gpu.device  # "gpu"
x_gpu.dtype   # 여전히 Float64(논리), 내부 MLX dtype은 float32

• 포인트: 사용자 시점에서는 Float64가 유지된다. MLX 내부는 float32지만, Numeric이 이를 감춰준다.

2️⃣ 예시: bit-free Float 자동 결정

data = np.array([1.0, 2.0], dtype=np.float64)
x = lucid.tensor(data, dtype=lucid.Float)  # bit-free
x.dtype   # Float64 (CPU)

y = x.to("gpu")
y.dtype   # Float (논리적으로 동일), GPU에서는 float32로 저장

• 포인트: 한 번도 Float32/Float64를 명시하지 않았지만, CPU에서는 64, GPU에서는 32로 알아서 맞춰진다.

3️⃣ 예시: one_hot dtype 지정

from lucid.nn.functional import one_hot
labels = lucid.tensor([0, 2, 1], dtype=lucid.Int64)
oh = one_hot(labels, num_classes=3, dtype=lucid.Float16)  # GPU라면 mx.float16

• Numeric.parse("gpu")가 mx.float16을 반환해 one_hot 결과 텐서가 원하는 dtype으로 바로 생성된다.

4️⃣ 예시: random.permutation과 dtype

perm = lucid.random.permutation(10, dtype=lucid.Int16, device="gpu")
perm.dtype   # Int16 (논리)
perm.device  # gpu, 내부 dtype은 mx.int16

• Random 경로도 Numeric을 통일적으로 사용해 CPU/GPU 차이를 숨긴다.

🧵 Numeric 도입 전후 코드 비교

# 도입 전: dtype을 직접 NumPy/MLX로 분기
if device == "cpu":
    data = np.array(..., dtype=np.float32)
else:
    data = mx.array(..., dtype=mx.float32)

# 도입 후: Numeric으로 한번에 표현
x = Tensor(data, dtype=lucid.Float32, device=device)

• 효과: 백엔드 분기 코드가 사라지고, 연산/모듈/랜덤 함수에서 dtype 파라미터를 통일할 수 있었다.

💡 Numeric과 디바이스 충돌 처리

• @func_op가 Tensor를 받으면 _check_is_tensor(..., device=device)로 승격하며, 고정된 텐서가 다른 디바이스면 RuntimeError를 던진다(lucid/_backend/core.py). dtype은 Numeric이 책임지므로, 디바이스 충돌과 dtype 충돌을 분리해서 생각할 수 있게 됐다.
• GPU 인덱싱에서 NumPy 배열을 사용하면 parse_mlx_indexing이 에러를 던진다. dtype 체계와 별개지만, GPU dtype을 강제하는 가드로서 Numeric의 장치 일관성 철학과 맞닿아 있다.

🧭 도입 결정 과정 회고

• 자체 타입의 목적: np.dtype과 mx.Dtype를 직접 노출하면 사용자가 "GPU면 mx.float32, CPU면 np.float32"처럼 조건문을 써야 했다.
• PyTorch 따라가기: torch.float32 같은 심플한 인터페이스를 제공하되, MLX 제약을 자동으로 처리하는 "PyTorch 같은데 Apple 실리콘 친화적인" 경험을 지향했다.
• 레거시 호환: 기존 코드에서 dtype=int 같은 builtin을 넘기던 패턴을 유지하기 위해 _dtype_map을 유지하고, Numeric과 섞여도 동작하도록 했다.

🧪 테스트/검증 전략

CPU/GPU에서 동일 스크립트를 돌려 dtype이 논리적으로 동일하게 보이는지 확인했고, float64 → float32 강등이 사용자 관점에서 드러나지 않는지 수동으로 체크했다. float16/float32 혼합에서도 grad 누락이 없는지 _match_grad_shape 동작을 살폈고, MLX lazy 환경에서는 loss.eval() 호출 전후 dtype 변동이 없는지 반복 확인했다.

⚔️ 남은 리스크와 TODO

• GPU float64 완전 지원: MLX가 향후 GPU float64를 지원하면 강등 로직을 조건부로 완화해야 한다. Numeric 생성자에서 디바이스 능력을 더 정교하게 검사하도록 개선 필요.
• bfloat16 지원: 현재 프리셋에 없다. MLX/NumPy 모두 지원 시 Numeric 프리셋을 확장해야 한다.
• dtype 추론 성능: auto_parse가 dtype 문자열을 파싱하므로 빈번히 호출되면 오버헤드가 생길 수 있다. 캐싱 레이어 검토.
• 컴파일러 친화성: JIT/그래프 최적화 시 Numeric → 실제 dtype 변환이 명시적으로 드러나야 한다. 향후 정적 그래프 변환을 염두에 두고 API를 더 명시적으로 할 수도 있다.

🔬 _dtype_bits와 auto_parse 파이프라인 심층

• 위치: lucid/types.py
• _dtype_bits는 NumPy dtype, MLX dtype, 문자열 모두에서 바이트 수를 추출한다. MLX에서는 dtype.size를 사용하고, NumPy/문자열은 np.dtype(...).itemsize로 처리한다.
• auto_parse는 이 값을 가져와 base_dtype.__name__ + bits 문자열을 만든 뒤 getattr(np | mx, new_dtype)로 실제 dtype을 생성한다. 덕분에 bit-free Numeric이 데이터 주도형으로 해석된다.

def auto_parse(self, data_dtype: type, device: _DeviceType) -> type | None:
    bits = self._dtype_bits(data_dtype)
    new_dtype = self.base_dtype.__name__ + str(bits)
    return getattr(np if device == "cpu" else mx, new_dtype, None)

• 설계 의도: dtype을 문자열 기반으로 생성해 NumPy/MLX 모두 같은 경로를 사용하게 했다. MLX가 새로운 dtype을 지원하면 자동으로 따라갈 수 있는 확장성을 확보했다.

🧭 구현 타임라인 메모

• 1단계 (CPU only): dtype은 순수 NumPy로만 처리, Numeric 필요성 미약.
• 2단계 (MLX 추가): GPU float64 문제를 맞닥뜨리며 Numeric 설계 착수. 강등 로직 추가.
• 3단계 (bit-free 도입): Float/Int/Complex 프리셋에 bits=None을 허용, auto_parse로 데이터 기반 결정.
• 4단계 (API 노출): lucid.__init__에 Numeric 프리셋 re-export, 기존 코드와 혼용할 수 있도록 _BuiltinNumeric 병행 지원.

🛑 실패 사례 모음 (빈번한 디버그 포인트)

• 복소수와 MLX: Complex64가 MLX에서 기대대로 동작하는지 확인이 필요했다. mx.complex64가 없던 시점에 오류가 발생해, 일단 Complex64만 프리셋으로 두고 더 높은 정밀도는 유보.
• dtype 문자열 파싱 실패: 일부 mx.float32가 float32로만 노출되어 regex가 정상 동작했지만, 사용자 커스텀 dtype 문자열이 들어오면 실패할 수 있었다. 명확한 예외 메시지(Unsupported dtype)를 던지도록 _dtype_bits에서 TypeError를 유도해 디버깅 시간을 줄였다.
• bool 혼용: Tensor.astype(bool) 호출 시 _is_bool_tensor 플래그를 갱신하지 않아 MLX에서 잘못된 캐스팅이 발생. bool이면 플래그와 dtype을 함께 업데이트하도록 수정.

🧭 실전 사용 패턴 모음

• 모델 정의 시

class MLP(lucid.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = lucid.nn.Linear(128, 64, dtype=lucid.Float32)
        self.fc2 = lucid.nn.Linear(64, 10, dtype=lucid.Float32)

• 데이터 로딩 후 GPU 이동

ds = TensorDataset(X, y, dtype=lucid.Float)  # bit-free
ds.to("gpu")  # GPU에서는 float32로 자동 배치

• 학습 루프

for x, y in loader:
    x = x.to("gpu")
    y = y.to("gpu")
    out = model(x)
    loss = criterion(out, y)
    loss.eval()
    loss.backward()
    opt.step()

• dtype 전환

x = lucid.randn(2, 2, dtype=lucid.Float16, device="gpu")
x32 = x.astype(lucid.Float32)  # MLX float32로 변환

🪄 Numeric이 Lucid 전체 구조에 준 의미

• API 일관성: Random, Tensor, Module, Functional 등 주요 영역에서 하나의 dtype 명명법을 사용하게 되었다.
• 디바이스 추상화: CPU/GPU를 넘나드는 코드에서 dtype 분기를 없앴다. 대신 디바이스 분기(.to)만 남겨 간결해졌다.

---

🌅 마무리와 앞으로의 방향

Numeric은 작은 클래스지만, 백엔드 이중화라는 근본 문제를 사용자에게서 숨겨주는 핵심 부품이다. 이 작업을 통해 다음과 같은 교훈을 얻었다.

• 디바이스 제약을 타입 시스템으로 승화하면, 사용자 코드에서 조건문이 사라지고, 오류를 더 일찍/명확히 잡을 수 있다.
• bit-free 타입처럼 "느슨한 추상"을 도입하면, 백엔드가 달라도 한 API를 공유할 수 있다.

앞으로는 bfloat16 등 추가 프리셋, float64 지원 확장, JIT 친화적인 dtype 선언을 탐색해볼 생각이다.