code: def test_custom_root_vector_nonlinear(self):
      def nonlinear_func(x, y):
          # func(x, y) == 0 if and only if x == y.
          return (x - y) * (x**2 + y**2 + 1)
  
      def tangent_solve(g, y):
          return jnp.linalg.solve(
              jax.jacobian(g)(y).reshape(-1, y.size), y.ravel()
          ).reshape(y.shape)
  
      def nonlinear_solve(y):
          f = lambda x: nonlinear_func(x, y)
          x0 = -jnp.ones_like(y)
          return lax.custom_root(f, x0, newton_raphson, tangent_solve)
  
      y = self.rng().randn(3, 1)
      jtu.check_grads(
          nonlinear_solve,
          (y,),
          order=2,
          rtol={jnp.float32: 1e-2, jnp.float64: 1e-3},
      )
  
      actual = jax.jit(nonlinear_solve)(y)
      self.assertAllClose(y, actual, rtol=1e-5, atol=1e-5)
  
imports: import re
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  import jax
  from jax import lax
  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 3 for calling the CustomRootTest obj's underlying member method code with example usage: obj.test_custom_root_vector_nonlinear() and returns: unknown, jnp, lax