import numpy as np

class SGD:
    def __init__(self, params, lr=0.01):
        self.params = params
        self.lr = lr
        
    def step(self):
        for layer in self.params:
            # Update Latent Weights
            # W_new = W_old - lr * grad
            layer.weight_latent -= self.lr * layer.grad_weight
            
            # Zero Gradients
            layer.grad_weight.fill(0)
            
    def zero_grad(self):
        for layer in self.params:
            layer.grad_weight.fill(0)
