code: def test_grad_norm(self):
      e = None
      try:
          with jax.debug_nans(True):
              jax.grad(jnp.linalg.norm)(jnp.zeros((3, 3), jnp.float32))
      except FloatingPointError as exc:
          e = exc
      self.assertIsNot(e, None)
      self.assertIn("invalid value", str(e))
      if sys.version_info >= (3, 11):
          self.assertIsInstance(
              e.__cause__,
              source_info_util.JaxStackTraceBeforeTransformation,
          )
      else:
          self.assertIsInstance(
              e.__cause__.__cause__,
              source_info_util.JaxStackTraceBeforeTransformation,
          )
  
imports: import re
  import sys
  import traceback
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  from jax import grad, jit, vmap, lax
  from jax._src import config
  from jax._src import core
  from jax._src import source_info_util
  from jax._src import test_util
  from jax._src import traceback_util
label_desc: Use image node_id 1 for calling the UserContextTracebackTest obj's underlying member method code with example usage: obj.test_grad_norm() without return types