code: def newton_raphson(func, x0):
      tol = 1e-16
      max_it = 20
  
      fx0, dfx0 = func(x0), jax.jacobian(func)(x0)
      initial_state = (
          0,
          x0,
          fx0,
          dfx0,
      )  # (iteration, x, f(x), grad(f)(x))
  
      def cond(state):
          it, _, fx, _ = state
          return (jnp.max(jnp.abs(fx)) > tol) & (it < max_it)
  
      def body(state):
          it, x, fx, dfx = state
          step = jnp.linalg.solve(
              dfx.reshape((-1, fx.size)), fx.ravel()
          ).reshape(fx.shape)
          x_next = x - step
          fx, dfx = func(x_next), jax.jacobian(func)(x_next)
          return (it + 1, x_next, fx, dfx)
  
      _, x, _, _ = lax.while_loop(cond, body, initial_state)
  
      return x
  
imports: from functools import partial
  import re
  import unittest
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  import jax
  from jax import lax
  from jax.ad_checkpoint import checkpoint
  from jax._src import test_util
  from jax import tree_util
  import jax.numpy
  import jax.scipy
  from jax import config
label_desc: Use image node_id 4 for calling a global function with example usage: newton_raphson(func, x0) and returns: x, unknown, unknown,  x_next,  fx,  dfx