Repository: tensorflow
Issue: Issue with softmax warning appearing in Tensorflow 2.17.0
Labels: type:bug,2.17

### TensorFlow version

2.17.0

### OS platform and distribution

Google Colab

### Current behavior?

Hi, everyone.

I am practicing implementing a Transformer model that machine translates English into Korean by reading TensorFlow guides and books. However, I am having trouble because an unknown UserWarning appears during the final translation process. In that issue, I've never used softmax before, but warning me about using it. This problem appears after the model has finished training and when making inferences.

I searched to see if there were any cases similar to mine, but it seems that no solution was found in any of them. 



same issue: https://github.com/tensorflow/tensorflow/issues/67758



This problem did not exist in tensorflow 2.15 but appeared in 2.17.0. I can't even guess what could be causing it. For those of you who are curious about the full code, I am leaving a Colab link. You can easily reproduce it by running it with Ctrl + F9 in Google Colab. The execution time of the entire code is approximately 5 minutes ~ 5 minutes and 30 seconds on a T4 GPU. That issue is at the bottom.

Colab Link: https://colab.research.google.com/drive/1IMFWoJ1s5ReKU9LYENROpAsZ47D6cG8T?usp=sharing

The data I used is 'kor-eng.zip' located at "https://www.manythings.org/anki/".

I'm really sorry for not writing the comments in English.

### Standalone code to reproduce the issue

```shell
class Transformer(keras.Model):
    def __init__(self, *, num_layers, encoder_sequence_length, decoder_sequence_length, source_vocab_size, target_vocab_size, embed_dim,
                 dense_dim, num_heads, dropout_rate):
        super().__init__()
        self.encoder = Encoder(num_layers=num_layers, sequence_length=encoder_sequence_length, input_dim=source_vocab_size, embed_dim=embed_dim,
                              dense_dim=dense_dim, num_heads=num_heads, dropout_rate=dropout_rate)
        self.decoder = Decoder(num_layers=num_layers, sequence_length=decoder_sequence_length, input_dim=target_vocab_size, embed_dim=embed_dim,
                              dense_dim=dense_dim, num_heads=num_heads, dropout_rate=dropout_rate)
        self.final_layer = tf.keras.layers.Dense(units=target_vocab_size)

    def call(self, inputs):
        encoder_inputs, decoder_inputs = inputs
        encoder_pad_mask = tf.math.not_equal(encoder_inputs, 0)[:, tf.newaxis]
        decoder_pad_mask = tf.math.not_equal(decoder_inputs, 0)[:, tf.newaxis]

        decoder_sequence_length = tf.shape(decoder_inputs)[1]
        causal_mask = tf.linalg.band_part(tf.ones((decoder_sequence_length, decoder_sequence_length), tf.bool), -1, 0)

        encoder_inputs = self.encoder(inputs=encoder_inputs, encoder_pad_mask=encoder_pad_mask)  # Shape: (batch_size, encoder_sequence_length, embed_dim)
        decoder_inputs = self.decoder(inputs=decoder_inputs, encoder_outputs=encoder_inputs,
                                      encoder_pad_mask=encoder_pad_mask, decoder_pad_mask=decoder_pad_mask,
                                      causal_mask=causal_mask)  # Shape: (batch_size, decoder_sequence_length, embed_dim)

        logits = self.final_layer(decoder_inputs)  # Shape: (batch_size, decoder_sequence_length, target_vocab_size)

        try:
            # losses/metrics가 커지지 않게 keras_mask를 제거
            del logits._keras_mask
        except AttributeError:
            pass

        return logits


class Translator(tf.Module):
    # 데이터 전처리 함수
    @staticmethod
    def preprocess_text(text_: str, max_repeat: int=2) -> str:
        """텍스트 문자열 중 일부 규칙적인 부분을 전처리 하는 함수

        Args:
            text_: 텍스트 문자열 -> str
            max_repeat: 똑같은 문자열이 연속해서 반복할 수 있는 최대 횟수 -> int

        Returns:
            text_: 전처리 된 텍스트 문자열 -> str
        """

        text_ = text_.lower()
        text_ = re.sub(pattern=rf"[^\w\s{string.punctuation}]", repl=r"", string=text_)
        text_ = re.sub(pattern=r"\?+", repl=r"?", string=text_) # ?가 2번 이상 연속되면 ?로 수정
        text_ = re.sub(pattern=r"\!+", repl=r"!", string=text_) # !가 2번 이상 연속되면 !로 수정
        text_ = re.sub(pattern=r"(?P<char>\D)(?P=char){" + str(max_repeat - 1) + r",}", repl=r"\g<char>" * max_repeat, string=text_) # 숫자가 아닌 똑같은 문자열이 repeat번 이상 연속되면 repeat만큼으로 수정
        text_ = re.sub(pattern=r"\.{2,}", repl=r"...", string=text_) # ..을 ...으로 변경
        text_ = re.sub(pattern=r"\.\.\.(?P
\w)", repl=r"... \g
", string=text_) # 띄어쓰기 교정
        text_ = re.sub(pattern=r"\s+", repl=" ", string=text_)

        # 영어 축약형을 풀기
        # 's 축약은 is / has 또는 아예 소유격으로 가능해서 제외
        # 'd 축약은 had / would / could 등으로 가능해서 제외
        # 'll 축약은 shall / will로 가능해서 제외
        text_ = re.sub(pattern=r"\bi'm\b", repl=r"i am", string=text_)
        text_ = re.sub(pattern=r"\b(?P<subj>you|we|they|there|who|when|where|what|how|why)'re\b", repl=r"\g<subj> are", string=text_)
        text_ = re.sub(pattern=r"\b(?P<verb>is|are|was|were|do|does|did|have|has|had|must|should|may|might|could|would|ought|dare|need)n't\b", repl=r"\g<verb> not", string=text_)
        text_ = re.sub(pattern=r"\bwon't\b", repl=r"will not", string=text_)
        text_ = re.sub(pattern=r"\bcan't\b", repl=r"can not", string=text_)
        text_ = re.sub(pattern=r"\bshan't\b", repl=r"shall not", string=text_)
        text_ = re.sub(pattern=r"\b(?P<subj>i|you|they|we|should|could|would|must|not)'ve\b", repl=r"\g<subj> have", string=text_)

        text_ = text_.strip()
        return text_

    # tensorflow의 바이트 문자열 데이터 전처리 함수
    @tf.function
    def tf_preprocess_text(text_: tf.string, max_repeat: int=2) -> tf.string:
        """텍스트 문자열 중 일부 규칙적인 부분을 전처리 하는 함수

        Args:
            text_: 바이트 문자열 -> tf.string
            max_repeat: 똑같은 문자열이 연속해서 반복할 수 있는 최대 횟수 -> int

        Returns:
            text_: 전처리 된 바이트 문자열 -> tf.string
        """

        text_ = tf.strings.lower(input=text_)
        text_ = tf.strings.regex_replace(input=text_, pattern=rf"[^\w\s{string.punctuation}]", rewrite=r"")
        text_ = tf.strings.regex_replace(input=text_, pattern=r"\?+", rewrite=r"?") # ?가 2번 이상 연속되면 ?로 수정
        text_ = tf.strings.regex_replace(input=text_, pattern=r"\!+", rewrite=r"!") # !가 2번 이상 연속되면 !로 수정

        # 숫자가 아닌 똑같은 문자열이 repeat번 이상 연속되면 repeat만큼으로 수정
        stacks = tf.TensorArray(dtype=tf.string, size=0, dynamic_size=True)
        for s in tf.strings.bytes_split(text_):
            if stacks.size() >= max_repeat:
                if tf.strings.regex_full_match(input=s, pattern=r"\D"):
                    back_s = stacks.gather(indices=tf.range(start=stacks.size() - max_repeat, limit=stacks.size(), delta=1))
                    if tf.math.reduce_all(back_s == s):
                        continue

            stacks = stacks.write(stacks.size(), s)

        text_ = tf.strings.reduce_join(inputs=stacks.stack())

        text_ = tf.strings.regex_replace(input=text_, pattern=r"\.{2,}", rewrite=r"...") # ..을 ...으로 변경
        text_ = tf.strings.regex_replace(input=text_, pattern=r"\.\.\.(\w)", rewrite=r"... \1") # 띄어쓰기 교정
        text_ = tf.strings.regex_replace(input=text_, pattern=r"\s+", rewrite=r" ")

        # 영어 축약형을 풀기
        # 's 축약은 is / has 또는 아예 소유격으로 가능해서 제외
        # 'd 축약은 had / would / could 등으로 가능해서 제외
        # 'll 축약은 shall / will로 가능해서 제외
        text_ = tf.strings.regex_replace(input=text_, pattern=r"\bi'm\b", rewrite=r"i am")
        text_ = tf.strings.regex_replace(input=text_, pattern=r"\b(you|we|they|there|who|when|where|what|how|why)'re\b", rewrite=r"\1 are")
        text_ = tf.strings.regex_replace(input=text_, pattern=r"\b(is|are|was|were|do|does|did|have|has|had|must|should|may|might|could|would|ought|dare|need)n't\b", rewrite=r"\1 not")
        text_ = tf.strings.regex_replace(input=text_, pattern=r"\bwon't\b", rewrite=r"will not")
        text_ = tf.strings.regex_replace(input=text_, pattern=r"\bcan't\b", rewrite=r"can not")
        text_ = tf.strings.regex_replace(input=text_, pattern=r"\bsha't\b", rewrite=r"shall not")
        text_ = tf.strings.regex_replace(input=text_, pattern=r"\b(i|you|they|we|should|could|would|must|not)'ve\b", rewrite=r"\1 have")

        text_ = tf.strings.strip(input=text_)
        return text_

    def __init__(self, source_tokenizer, target_tokenizer, target_length, model):
        super().__init__()
        self.source_tokenizer = source_tokenizer
        self.target_tokenizer = target_tokenizer
        self.target_length = target_length
        self.model = model

    def __call__(self, sentence):
        sentence = Translator.tf_preprocess_text(text_=sentence)
        sentence_token = self.source_tokenizer.tokenize(sentence)[tf.newaxis]

        encoder_input = sentence_token

        starts = tf.constant(2, dtype=tf.int32)[tf.newaxis]
        ends = tf.constant(3, dtype=tf.int32)[tf.newaxis]
        decoder_token = tf.TensorArray(dtype=tf.int32, size=0, dynamic_size=True)
        decoder_token = decoder_token.write(0, starts)

        for i in tf.range(self.target_length):
            decoder_input = tf.transpose(decoder_token.stack())
            predictions = self.model([encoder_input, decoder_input], training=False)
            predictions = predictions[0, -1, :] # 마지막 토큰으로부터 예측
            predicted_id = tf.argmax(input=predictions, output_type=tf.int32)[tf.newaxis]

            if predicted_id == ends:
                break

            decoder_token = decoder_token.write(i + 1, predicted_id)

        decoder_token = tf.transpose(decoder_token.stack())[0]
        decoder_token = decoder_token[1:]

        return self.target_tokenizer.detokenize(decoder_token)
```


### Relevant log output

```shell
/usr/local/lib/python3.10/dist-packages/keras/src/ops/nn.py:545: UserWarning: You are using a softmax over axis 3 of a tensor of shape (1, 8, 1, 1). This axis has size 1. The softmax operation will always return the value 1, which is likely not what you intended. Did you mean to use a sigmoid instead?
  warnings.warn(
English: tom came here to learn french.
Translated Korean: 톰은 프랑스어를 배우러 여기에 왔어. 
Real Korean: 
 톰은 프랑스어를 배우려고 여기에 왔어. 

```