code: def class_reweighted_grad(
      logits,
      class_weights,
      norm_grads=True,
      name="class_reweighted_grad",
  ):
      def custom_grad_fn(op, g):
          new_g = g * class_weights
          if norm_grads:
              ratio = tf.math.divide_no_nan(tf.norm(g), tf.norm(new_g))
          else:
              ratio = 1.0
          return [new_g * ratio]
  
      @function.Defun(
          logits.dtype,
          func_name=name,
          python_grad_func=custom_grad_fn,
          shape_func=lambda _: [logits.get_shape()],
      )
      def identity(l):
          return tf.identity(l)
  
      return identity(logits)
  
imports: import functools
  import tensorflow
  from tensorflow_addons.text.crf import crf_log_likelihood
  from finetune.base_models.gpt.featurizer import attn, dropout, norm
  from finetune.util.shapes import shape_list, merge_leading_dims
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.errors import FinetuneError
  from finetune.nn.activations import act_fns
  from finetune.nn.nn_utils import norm
  from tensorflow.python.framework import function
label_desc: Use image node_id 11 for calling a global function with example usage: class_reweighted_grad(logits, class_weights, norm_grads, name) and returns: identity, list, tf