code: def _add_diagonal_regularizer(
      A: jnp.ndarray, diag_reg: float, diag_reg_absolute_scale: bool
  ) -> jnp.ndarray:
      dimension = A.shape[0]
      if not diag_reg_absolute_scale:
          diag_reg *= jnp.trace(A) / dimension
      return A + diag_reg * jnp.eye(dimension)
  
imports: import collections
  from functools import lru_cache
  from typing import Any, Callable, Generator, Iterable, NamedTuple, Optional, Protocol, Union
  import jax
  from jax import grad
  from jax.experimental import ode
  import jax.numpy
  import jax.scipy
  from jax.tree_util import tree_all
  from jax.tree_util import tree_map
  import numpy
  import scipy
  from .utils import dataclasses
  from .utils import utils
  from .utils.typing import Axes
  from .utils.typing import Get
  from .utils.typing import KernelFn
label_desc: Use image node_id 9 for calling a global function with example usage: _add_diagonal_regularizer(A, diag_reg, diag_reg_absolute_scale) and returns: unknown