code: def docrep_pos_embed(
      input_context, positional_channels, batch_size, seq_length, width
  ):
      """Embed 2D position DocRep-style, i.e. separate sinusoidal embeddings for each dim"""
      init = tf.compat.v1.variance_scaling_initializer(
          scale=0.02, mode="fan_avg", distribution="truncated_normal"
      )
      embedded_input_context = embed_position(
          input_context, positional_channels, batch_size, seq_length
      )
      return tf.compat.v1.layers.dense(
          embedded_input_context,
          width,
          use_bias=False,
          kernel_initializer=init,
      )
  
imports: import copy
  import json
  import math
  import numpy
  import tensorflow
  import functools
  from transformers.activations_tf import gelu
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.nn.auxiliary import embed_position
label_desc: Use image node_id 8 for calling a global function with example usage: docrep_pos_embed(input_context, positional_channels, batch_size, seq_length, width) and returns: tf