Jax¶

# from scipy.optimize import minimize

from jax.scipy.optimize import minimize
import jax.numpy as np

import jax.numpy as np
from jax import jit, value_and_grad

def costFunction(weights):
   # reshapes flattened weights into 2d matrix
   weights = np.reshape(weights, newshape=(100, 75))

   # weighted row-wise sum
   weighted = np.sum(x * weights, axis=1)

   # squared residuals
   residualsSquared = (y - weighted) ** 2

   return np.sum(residualsSquared)

# create the derivatives
obj_and_grad = jit(value_and_grad(costFunction))

minimize(obj_and_grad, x0=startingWeights.flatten(), jac=True)

import 
o.minimize(fun, x0, args=(), *, method, tol=None, options=None)