NumPy’den JAX’e Geçişte Karşılaştığım 3 Kritik Gerçek

NumPy ve PyTorch’un sunduğu standart veri bilimi araçları, projeler büyüdükçe performans sınırlamalarına yol açabiliyor. Bu nedenle, ben de 100 Days of AI Agents adlı açık kaynak projeme hız kazandırmak için JAX’e yöneldim. Ancak bu geçiş, sadece bir framework değişikliğinden ibaret değildi; tamamen yeni bir bakış açısı gerektiren bir paradigma kaymasıydı.

JAX’in sunduğu avantajları anlamak için birkaç gün boyunca temel mekanizmalarını inceledim. İşte karşılaştığım en önemli üç dönüm noktası ve bu süreçte beni etkileyen kod parçacıkları. Bu deneyimler, JAX’in neden modern yapay zeka geliştiricileri için vazgeçilmez bir araç olduğunu gösteriyor.

Dizinlere Erişimde Sabitlik: JAX’in Temel Taşı

NumPy’de alışık olduğumuz dizilere doğrudan erişim ve değiştirme işlemleri, JAX’te farklı bir yaklaşım gerektiriyor. Örneğin, NumPy’de basit bir dizi oluşturup ilk elemanını değiştirmek oldukça basitti:

import numpy as np

dizi = np.arange(10)
dizi[0] = 10
print(dizi)  # Çıktı: [10 1 2 3 4 5 6 7 8 9]

Ancak JAX’te aynı işlem yapıldığında tamamen farklı bir davranışla karşılaşıyoruz. JAX dizileri (jax.Array) oluşturulduktan sonra değiştirilemez. Bu durum, ilk bakışta kısıtlama gibi görünse de aslında JAX’in işlevsel programlama yaklaşımının temelinde yatan bir tasarım tercihi:

import jax.numpy as jnp

dizi = jnp.arange(10)
yeni_dizi = dizi.at[0].set(10)

print(yeni_dizi)  # Çıktı: [10 1 2 3 4 5 6 7 8 9]
print(dizi)       # Çıktı: [0 1 2 3 4 5 6 7 8 9]

Bu yöntem, bellekte yeni bir dizi oluştururken orijinal dizinin korunmasını sağlıyor. Bu durum, özellikle dağıtılmış sistemlerde yan etkilerin önlenmesi açısından kritik bir avantaj sunuyor. Tabii ki, kopyalama işlemi bellek kullanımını artırabilir, ancak uzun vadede güvenilirliği ve tutarlılığı artırıyor.

Donanım Farkındalığı: JAX’in En Büyük Gücü

JAX’in en etkileyici özelliklerinden biri, dizilerin nerede saklandığını otomatik olarak algılaması. NumPy’de GPU ya da TPU gibi donanımlara elle yönlendirme yapmanız gerekirken, JAX dizileri varsayılan olarak en hızlı erişilebilir donanıma otomatik olarak yönlendiriliyor.

Örneğin, yerel bilgisayarımda çalıştırdığım bir JAX dizisinin konumunu sorguladığımda:

import jax

dizi = jax.numpy.arange(10)
print(dizi.devices())

Çıktı olarak sadece CpuDevice(id=0) görüyorum. Eğer sistemde bir GPU bulunuyorsa, JAX dizileri otomatik olarak GPU’ya yönlendirilir. Bu özellik, özellikle çoklu cihazlı ortamlarda büyük bir kolaylık sağlıyor.

Ayrıca, JAX dizileri birden fazla cihaz arasında parçalanabilir. Parçalama durumunu aşağıdaki gibi kontrol edebilirsiniz:

print(dizi.sharding)

Çıktı olarak SingleDeviceSharding(device=CpuDevice(id=0), memory_kind=device) alıyorsunuz. Bu, JAX’in modern donanım ölçekleme için tasarlandığını açıkça gösteriyor. Sistem kaynaklarını en verimli şekilde kullanarak performansı artırmak, JAX’in en büyük avantajlarından biri.

JIT Derlemesiyle Performans Patlaması

JAX’in sunduğu en güçlü özelliklerden biri, Just-In-Time (JIT) derlemesi. JAX’in varsayılan davranışı, Python’un standart yürütme modeline benzer şekilde komutları teker teker çalıştırmaktır. Ancak fonksiyonları jax.jit dekoratörüyle sarmaladığınızda, JAX tüm işlemleri tek seferde optimize eder ve birlikte çalıştırır.

Bu avantajı test etmek için basit bir normalizasyon fonksiyonu yazdım ve performansını ölçtüm:

from jax import jit
import jax.numpy as jnp
import numpy as np

def normalize_veri(X):
    X = X - X.mean(0)
    return X / X.std(0)

normalize_derlenmis = jit(normalize_veri)

# Test verisi oluşturma
np.random.seed(22)
veri = jnp.array(np.random.rand(100000, 10))

Performansı ölçmek için %timeit komutunu kullanarak hem normal hem de derlenmiş fonksiyonları karşılaştırdım. Sonuçlar oldukça çarpıcıydı:

• Standart yürütme: 1.52 ms ± 16.3 μs
• JIT derlemesiyle yürütme: 1.16 ms ± 26.2 μs

JIT derlemesi, işlemlerin önceden bilinir olması sayesinde önemli bir hız artışı sağlıyor. Ancak unutulmaması gereken bir nokta var: Tüm JAX kodları JIT ile derlenemez. Derleme için dizilerin boyutlarının derleme zamanında sabit ve biliniyor olması gerekiyor.

Sonraki Adımlar: JAX’in Derinliklerine Dalmak

Bu deneyimler, JAX’in sunduğu potansiyelin sadece yüzeyini çiziyor. Gelecekte derinlemesine inceleyeceğim konular arasında fonksiyonel rastgelelik (jax.random), otomatik türev alma (jax.grad) ve otomatik vektörizasyon (jax.vmap) bulunuyor. Bu araçlar, yapay zeka projelerini hem daha hızlı hem de daha güvenilir hale getiriyor.

Peki siz? JAX’e geçiş yaptınız mı? En büyük zorluklarınız nelerdi? Deneyimlerinizi aşağıdaki yorumlarda paylaşın ve bu heyecan verici araç hakkında daha fazla bilgi edinmek isteyenlere yol gösterin.