code: def additive_weight_decay(
      weight_decay: float = 0.0,
  ) -> GradientTransformation:
      """Add parameter scaled by `weight_decay`, to all parameters with more than one dim (i.e. exclude ln, bias etc)
  
      Args:
        weight_decay: a scalar weight decay rate.
  
      Returns:
        An (init_fn, update_fn) tuple.
      """
  
      def init_fn(_):
          return AdditiveWeightDecayState()
  
      def update_fn(updates, state, params):
          updates = jax.tree_multimap(
              lambda g, p: g + weight_decay * p * (len(g.shape) > 1),
              updates,
              params,
          )
          return updates, state
  
      return GradientTransformation(init_fn, update_fn)
  
imports: import jax
  import jax.numpy
  from jax.experimental.pjit import with_sharding_constraint
  from optax import AdditiveWeightDecayState, GradientTransformation, OptState
label_desc: Use image node_id 5 for calling a global function with example usage: additive_weight_decay(weight_decay) and returns: GradientTransformation, AdditiveWeightDecayState, updates,  state