code: def safe_norm(x, axis=-1, keepdims=False, tol=1e-9):
      """Calculates a np.linalg.norm(d) that's safe for gradients at d=0.
  
      These gymnastics are to avoid a poorly defined gradient for np.linal.norm(0)
      see https://github.com/google/jax/issues/3058 for details
  
  
      Args:
        x: A np.array
        axis: The axis along which to compute the norm
        keepdims: if True don't squeeze the axis.
        tol: the absolute threshold within which to zero out the gradient.
  
      Returns:
        Equivalent to np.linalg.norm(d)
      """
      return jnp.linalg.norm(x, axis=axis, keepdims=keepdims)
  
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 1 for calling a global function with example usage: safe_norm(x, axis, keepdims, tol) and returns: jnp