code: def minimize(f, x, num_steps=10000, step_size=0.000001, mass=0.9):
      opt_init, opt_update, get_params = optimizers.momentum(
          step_size, mass
      )
  
      @jit
      def update(i, opt_state):
          x = get_params(opt_state)
          return opt_update(i, grad(f)(x), opt_state)
  
      opt_state = opt_init(x)
      for i in range(num_steps):
          opt_state = update(i, opt_state)
      return get_params(opt_state)
  
file_path: jax/examples/kernel_lsq.py
imports: from functools import partial
  import numpy.random
  import jax.numpy
  from jax.example_libraries import optimizers
  from jax import grad, jit, make_jaxpr, vmap, lax
label_desc: Use image node_id 2 for calling a global function with example usage: minimize(f, x, num_steps, step_size, mass) and returns: get_params, opt_update