Levenberg-Marquardt in tensorflow

Question

I wish to find out if anyone has tried to implement the well-known Levenberg-Marquardt algorithm in tensorflow? I have a number of issues on trying to implement it, during parameter updates. The following code snippet shows an implementation of the update function:

def func_var_update(cost, parameters):
  # compute gradients or Jacobians for cost with respect to parameters
  dloss_dw = tf.gradients(cost, parameters)[0]
  # Return dimension of gradient vector
  dim, _ = dloss_dw.get_shape()
  # Compute hessian matrix using results of gradients
  hess = []
  for i in range(dim):
        # Compute gradient ot Jacobian matrix for loss function
        dfx_i = tf.slice(dloss_dw, begin=[i,0] , size=[1,1])
        ddfx_i = tf.gradients(dfx_i, parameters)[0] 
        # Get the actual tensors at the end of tf.gradients
        hess.append(ddfx_i)
    hess = tf.squeeze(hess)     
    dfw_new = tf.diag(dloss_dw)
    # Update factor consisting of the hessian, product of identity matrix and Jacobian vector
    JtJ = tf.linalg.inv(tf.ones((parameters.shape[0], parameters.shape[0])) + hess)
    # product of gradient and damping parameter
    pdt_JtJ = tf.matmul(JtJ, dloss_dw)
    # Performing update here
    new_params = tf.assign(parameters, parameters - pdt_JtJ)
  return new_params

And the following call:

def mainfunc()
    with tf.Session():
   .....
       vec_up = sess.run(func_var_update(), feed_dict=....)

results in the following error:

InvalidArgumentError (see above for traceback): Input is not invertible.

But Both the dimension of the Jacobian/gradient and hessian are okay when I print them during runtime. The other problem I have is not being able to keep track of the parameters after each update and then adapt them to personal needs before feeding them into the optimizer later. I wanted to fix some parameters, and compute hessian and jacobian for others while performing optimization at the same time. Any help will be appreciated.