b3377b098a78e8a2508527f9d71a89b15fef3929,official/nlp/modeling/layers/on_device_embedding.py,OnDeviceEmbedding,call,#OnDeviceEmbedding#Any#,80

Before Change

  def call(self, inputs):
    input_shape = tf_utils.get_shape_list(inputs, expected_rank=2)
    input_shape.append(self._embedding_width)
    flat_inputs = tf.reshape(inputs, [-1])
    if self._use_one_hot:
      one_hot_data = tf.one_hot(
          flat_inputs, depth=self._vocab_size, dtype=self._dtype)

After Change

        embeddings,
        // Work around b/142213824: prefer concat to shape over a Python list.
        tf.concat([tf.shape(inputs), [self._embedding_width]], axis=0))
    embeddings.set_shape(inputs.shape.as_list() + [self._embedding_width])
    return embeddings

Instances