code: def compute_max_norm(tensors: PyTree) -> float:
      """Compute the maximum norm in a pytree of tensors."""
      leaves, _ = jax.tree_util.tree_flatten(tensors)
      norms = jnp.sqrt(sum(jnp.vdot(x, x) for x in leaves))
      max_norm = jnp.max(norms)
      return max_norm
  
file_path: scenic/scenic/projects/objectvivit/train_utils.py
imports: from typing import Any, Dict, Tuple
  from absl import logging
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 6 for calling a global function with example usage: compute_max_norm(tensors) and returns: max_norm