code: def featurizer(
      X,
      encoder,
      config,
      train=False,
      reuse=None,
      encoder_state=None,
      context=None,
      context_dim=None,
      **kwargs,
  ):
      """
      The main element of the OSCAR model. Maps from tokens ids to a dense, embedding of the sequence.
  
      :param X: A tensor of token indexes with shape [batch_size, sequence_length, token_idx]
      :param encoder: A TextEncoder object.
      :param config: A config object, containing all parameters for the featurizer.
      :param train: If this flag is true, dropout and losses are added to the graph.
      :param reuse: Should reuse be set within this scope.
      :return: A dict containing;
          embed_weights: the word embedding matrix.
          features: The output of the featurizer_final state.
          sequence_features: The output of the featurizer at each timestep.
      """
      initial_shape = [a or -1 for a in X.get_shape().as_list()]
      if len(initial_shape) != 2:
          X = tf.reshape(X, shape=[-1] + initial_shape[-1:])
  
      x_shape = tf.shape(input=X)
      with tf.compat.v1.variable_scope("model/featurizer", reuse=reuse):
          clf_token = encoder.end_token
          pool_idx = tf.cast(
              tf.argmax(
                  input=tf.cast(tf.equal(X, clf_token), tf.float32),
                  axis=1,
              ),
              tf.int32,
          )
          if encoder_state is None:
              embed_weights = tf.compat.v1.get_variable(
                  "we",
                  [
                      encoder.vocab_size + config.max_length,
                      config.n_embed,
                  ],
                  initializer=tf.compat.v1.random_normal_initializer(
                      stddev=config.weight_stddev
                  ),
              )
          else:
              embed_weights = encoder_state["embed_weights"]
  
          if config.oscar_use_fp16:
              embed_weights = tf.cast(embed_weights, tf.float16)
  
          if config.train_embeddings:
              embed_weights = dropout(
                  embed_weights, config.embed_p_drop, train
              )
          else:
              embed_weights = tf.stop_gradient(embed_weights)
  
          X = tf.reshape(X, [-1, x_shape[1]])
          h = tf.gather(embed_weights, X)
          if config.oscar_use_timing:
              pos_values = get_pos_values(
                  x_shape[1], encoder.vocab_size
              )
              h += tf.gather(embed_weights, pos_values)
  
          for layer in range(config.n_layer):
              with tf.compat.v1.variable_scope("h%d_" % layer):
                  if (
                      (config.n_layer - layer)
                      == config.num_layers_trained
                      and config.num_layers_trained != config.n_layer
                  ):
                      h = tf.stop_gradient(h)
  
                  block_fn_fwd = functools.partial(
                      block,
                      block_name="block%d_" % layer,
                      use_fp16=config.oscar_use_fp16,
                      pool_idx=None,
                      encoder_state=encoder_state,
                      train=train,
                      pdrop=config.resid_p_drop,
                      use_fused_kernel=config.oscar_use_fused_kernel,
                  )
  
                  if config.low_memory_mode and train:
                      block_fn_fwd = recompute_grad(
                          block_fn_fwd, use_entire_scope=True
                      )
                  h = block_fn_fwd(h)
  
          h = normal_1d_conv_block(
              h, 1, "output", config.oscar_use_fp16, dilation=1
          )
  
          mask = tf.expand_dims(
              tf.sequence_mask(
                  pool_idx, maxlen=tf.shape(input=h)[1], dtype=h.dtype
              ),
              -1,
          )
  
          if config.oscar_feat_mode == "clf_tok":
              clf_h = tf.gather_nd(
                  h, tf.stack([tf.range(shape_list(h)[0]), pool_idx], 1)
              )
          elif config.oscar_feat_mode == "mean_tok":
              clf_h = tf.reduce_sum(
                  input_tensor=h * mask, axis=1
              ) / tf.reduce_sum(input_tensor=h)
          elif config.oscar_feat_mode == "max_tok":
              clf_h = tf.reduce_max(
                  input_tensor=h - (1e5 * (1.0 - mask)), axis=1
              )
          else:
              raise ValueError(
                  "config.feat_mode should be one of clf_tok, mean_tok or max_tok"
              )
  
          if len(initial_shape) != 2:
              seq_feats = tf.reshape(
                  h, shape=initial_shape + [config.n_embed]
              )
          else:
              seq_feats = h
  
          return {
              "embed_weights": embed_weights,
              "features": cast_maybe(clf_h, tf.float32),
              "sequence_features": seq_feats,
              "eos_idx": pool_idx,
              "encoded_input": X[
                  :, : tf.reduce_min(input_tensor=pool_idx)
              ],
              "lengths": lengths_from_eos_idx(
                  eos_idx=pool_idx, max_length=shape_list(X)[1]
              ),
          }
  
file_path: finetune/finetune/base_models/oscar/featurizer.py
imports: import tensorflow
  from finetune.base_models.gpt.featurizer import dropout, embed, split_heads, merge_heads
  from finetune.util.shapes import shape_list, lengths_from_eos_idx
  from finetune.custom_ops import recursive_agg
  from finetune.base_models.gpt.featurizer import get_pos_values
  from finetune.optimizers.recompute_grads import recompute_grad
  import functools
label_desc: Use image node_id 11 for calling a global function with example usage: featurizer(X, encoder, config, train, reuse, encoder_state, context, context_dim) and returns: dict