TF: embeddings out of bounds check factored into function (#23427)
This commit is contained in:
@@ -53,7 +53,7 @@ from ...modeling_tf_utils import (
|
||||
keras_serializable,
|
||||
unpack_inputs,
|
||||
)
|
||||
from ...tf_utils import shape_list, stable_softmax
|
||||
from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
|
||||
from ...utils import logging
|
||||
from .configuration_{{cookiecutter.lowercase_modelname}} import {{cookiecutter.camelcase_modelname}}Config
|
||||
|
||||
@@ -126,16 +126,7 @@ class TF{{cookiecutter.camelcase_modelname}}Embeddings(tf.keras.layers.Layer):
|
||||
assert not (input_ids is None and inputs_embeds is None)
|
||||
|
||||
if input_ids is not None:
|
||||
# Note: tf.gather, on which the embedding layer is based, won't check positive out of bound
|
||||
# indices on GPU, returning zeros instead. This is a dangerous silent behavior.
|
||||
tf.debugging.assert_less(
|
||||
input_ids,
|
||||
tf.cast(self.vocab_size, dtype=input_ids.dtype),
|
||||
message=(
|
||||
"input_ids must be smaller than the embedding layer's input dimension (got"
|
||||
f" {tf.math.reduce_max(input_ids)} >= {self.vocab_size})"
|
||||
),
|
||||
)
|
||||
check_embeddings_within_bounds(input_ids, self.vocab_size)
|
||||
inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
|
||||
|
||||
input_shape = shape_list(inputs_embeds)[:-1]
|
||||
@@ -1670,7 +1661,7 @@ from ...modeling_tf_utils import (
|
||||
keras_serializable,
|
||||
unpack_inputs,
|
||||
)
|
||||
from ...tf_utils import shape_list, stable_softmax
|
||||
from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
|
||||
from ...utils import ContextManagers, logging
|
||||
from .configuration_{{cookiecutter.lowercase_modelname}} import {{cookiecutter.camelcase_modelname}}Config
|
||||
|
||||
@@ -2311,16 +2302,7 @@ class TF{{cookiecutter.camelcase_modelname}}Encoder(tf.keras.layers.Layer):
|
||||
if hasattr(self.embed_tokens, "load_weight_prefix"):
|
||||
context.append(tf.name_scope(self.embed_tokens.load_weight_prefix + "/"))
|
||||
with ContextManagers(context):
|
||||
# Note: tf.gather, on which the embedding layer is based, won't check positive out of bound
|
||||
# indices on GPU, returning zeros instead. This is a dangerous silent behavior.
|
||||
tf.debugging.assert_less(
|
||||
input_ids,
|
||||
tf.cast(self.embed_tokens.input_dim, dtype=input_ids.dtype),
|
||||
message=(
|
||||
"input_ids must be smaller than the embedding layer's input dimension (got"
|
||||
f" {tf.math.reduce_max(input_ids)} >= {self.embed_tokens.input_dim})"
|
||||
),
|
||||
)
|
||||
check_embeddings_within_bounds(input_ids, self.embed_tokens.input_dim)
|
||||
inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
|
||||
|
||||
embed_pos = self.embed_positions(input_shape)
|
||||
@@ -2518,16 +2500,7 @@ class TF{{cookiecutter.camelcase_modelname}}Decoder(tf.keras.layers.Layer):
|
||||
if hasattr(self.embed_tokens, "load_weight_prefix"):
|
||||
context.append(tf.name_scope(self.embed_tokens.load_weight_prefix + "/"))
|
||||
with ContextManagers(context):
|
||||
# Note: tf.gather, on which the embedding layer is based, won't check positive out of bound
|
||||
# indices on GPU, returning zeros instead. This is a dangerous silent behavior.
|
||||
tf.debugging.assert_less(
|
||||
input_ids,
|
||||
tf.cast(self.embed_tokens.input_dim, dtype=input_ids.dtype),
|
||||
message=(
|
||||
"input_ids must be smaller than the embedding layer's input dimension (got"
|
||||
f" {tf.math.reduce_max(input_ids)} >= {self.embed_tokens.input_dim})"
|
||||
),
|
||||
)
|
||||
check_embeddings_within_bounds(input_ids, self.embed_tokens.input_dim)
|
||||
inputs_embeds = self.embed_tokens(input_ids)
|
||||
|
||||
hidden_states = inputs_embeds
|
||||
|
||||
Reference in New Issue
Block a user