code: def _safe_norm_jvp(axis, keepdims, tol, primals, tangents):
      """Custom JVP rule for safe_norm."""
      (x,) = primals
      (x_dot,) = tangents
      safe_tol = max(tol, 1e-30)
      y = safe_norm(x, tol=safe_tol, axis=axis, keepdims=True)
      y_safe = jnp.maximum(y, tol)  # Prevent divide by zero.
      y_dot = jnp.where(
          y > safe_tol, x_dot * x / y_safe, jnp.zeros_like(x)
      )
      y_dot = jnp.sum(y_dot, axis=axis, keepdims=True)
      # Squeeze the axis if `keepdims` is True.
      if not keepdims:
          y = jnp.squeeze(y, axis=axis)
          y_dot = jnp.squeeze(y_dot, axis=axis)
      return y, y_dot
  
imports: import collections
  from concurrent import futures
  import contextlib
  import functools
  import time
  from typing import List, Union
  import jax
  from jax import tree_util
  import jax.numpy
  import numpy
  from scipy import interpolate
  from scipy.spatial import transform
  import tqdm
label_desc: Use image node_id 2 for calling a global function with example usage: _safe_norm_jvp(axis, keepdims, tol, primals, tangents) and returns: y,  y_dot