# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import numpy as onp
import jax.numpy as jnp
from jax import grad, jit, vmap, value_and_grad
import jax

# Generate key which is used to generate random numbers
init_rng_key = jax.random.PRNGKey(1)




# +
n_samples, n_dims = 100, 20

x = jnp.array(onp.random.randn(n_samples, n_dims))
w_true = jnp.array(onp.random.randn(n_dims))
y_true = jnp.dot(x, w_true)
print(x.shape)
print(w_true.shape)
print(y_true.shape)
w_init = jnp.array(onp.random.randn(n_dims))
print(w.shape)
true_params = {"w": w_true}
params = {"w": w_init}

print(params)
# -

x

w_true

y_true


# +
def forward(x, w):
    return jnp.dot(x, w) 

def mse(w, x, y_true):
    diff = y_true - jnp.dot(x, w)  
    return jnp.mean(diff**2)


def full_pass(w, x, y_true, lr):
    dloss = jax.grad(mse)
    dw = dloss(w, x, y_true)

    return w - lr * dw
   
def opt(w, x, y_true, lr=0.1, iters=100):
    for i in range(iters):
        old_w = w
        w = full_pass(w, x, y_true, lr=lr)
        y_pred = forward(x, w)
        print("diff in outputs", ((y_pred-y_true) ** 2).mean())
        print("diff in weights", ((w - old_w) ** 2).mean())
    return params



# -

new_w = opt(w, x, y_true)


# +
def mse(w, x, y_true):
    return jnp.mean((y_true - jnp.dot(x, w)) ** 2)

dloss = jax.grad(mse)
lr = 0.1
for i in range(100):
    dw = dloss(w, x, y_true)
    w -= lr * dw
    
# -

w - w_true

# +
w = jnp.array(onp.random.randn(n_dims))

y_true = forward(x, w_true)

def forward(x, w):
    return jnp.tanh(jnp.dot(x, w))

def mse(w, x, y_true):
    return jnp.mean((y_true - forward(x,w)) ** 2)

dloss = jax.grad(mse)
lr = 0.5
for i in range(100):
    dw = dloss(w, x, y_true)
    w -= lr * dw
    print( ((y_true - forward(x, w))**2).mean())

# +
# jax.grad??
# -



StarcoderdataJupyterScriptsDedupFiltered