code: def _get_weights(model_hparams, vocab_size, hidden_dim=None):
      """Copied from tensor2tensor/layers/modalities.py but uses total vocab."""
      if hidden_dim is None:
          hidden_dim = model_hparams.hidden_size
      num_shards = model_hparams.symbol_modality_num_shards
      shards = []
      for i in range(num_shards):
          shard_size = (sum(vocab_size) // num_shards) + (
              1 if i < sum(vocab_size) % num_shards else 0
          )
          var_name = "weights_%d" % i
          shards.append(
              tf.get_variable(
                  var_name,
                  [shard_size, hidden_dim],
                  initializer=tf.random_normal_initializer(
                      0.0, hidden_dim**-0.5
                  ),
              )
          )
      if num_shards == 1:
          ret = shards[0]
      else:
          ret = tf.concat(shards, 0)
      # Convert ret to tensor.
      if not tf.executing_eagerly():
          ret = common_layers.convert_gradient_to_tensor(ret)
      return ret
  
file_path: magenta/magenta/models/score2perf/modalities.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
label_desc: Use image node_id 1 for calling a global function with example usage: _get_weights(model_hparams, vocab_size, hidden_dim) and returns: ret