JAX در برابر PyTorch

model = MNISTModel().to(device)                  # 784, 256, 128, 10
criterion = nn.CrossEntropyLoss()               # لاجیت‌های خام را دریافت می‌کند
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(10):
    model.train()                               # دراپ‌اوت فعال است
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()                   # گرادیان‌ها انباشته می‌شوند
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()                         # نوار محاسباتی را بازپخش می‌کند
        optimizer.step()
torch.save(model.state_dict(), "mnist_mlp.pt")  # ~97.8٪ دقت آزمون

optimizer = optax.adam(learning_rate=1e-3)
@jax.jit                                                             # کامپایل به XLA
def train_step(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss                                   # خروجی وضعیت جدید
key = random.PRNGKey(0)
params = init_params(key)                                            # پای‌تری با دیکشنری‌های تودرتو
opt_state = optimizer.init(params)
for epoch in range(n_epochs):
    key, subkey = random.split(key)                                  # کلید جدید در هر ایپاک
    perm = random.permutation(subkey, len(X_train))
    params, opt_state, loss = train_step(params, opt_state, xb, yb)  # اندازه دسته: 128