# AUTOGENERATED! DO NOT EDIT! File to edit: 02_NN.ipynb (unless otherwise specified).

__all__ = ['layer_init', 'linear']

# Cell
import torch
from torch import Tensor

# Cell
def layer_init(h,w,c=None):
    return

# Cell
class linear:
    def __init__(self,in_ch,out_ch):
        self.weight = layer_init(in_ch,out_ch)
        self.shape = (in_ch,out_ch)
        # save for backward
        self.fpass = None
        self.grad = torch.zero_(in_ch,out_ch)

    def forward(self,x):
        out = x @ self.weight
        self.fpass = out
        return out

    def backward(self,bpass):
        self.grad = self.fpass.T @ bpass
        bpass = bpass @ (self.weight.T)
