code: def compute_max_norm(tensors: train_utils.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
  
imports: import functools
  from typing import Any, Callable, Dict, Iterator, Mapping, Optional, Tuple, Type, Union
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  import flax
  from flax import jax_utils
  import flax.linen
  import jax
  from jax.example_libraries.optimizers import clip_grads
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.av_mae import base_model
  from scenic.projects.av_mae import train_utils
  from scenic.train_lib_deprecated import lr_schedules
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import pretrain_utils
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 2 for calling a global function with example usage: compute_max_norm(tensors) and returns: max_norm