!pip install keras-nlp --upgrade --quiet

import tensorflow as tf
import keras_nlp

# Based on test_valid_call()
# https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/layers/modeling/masked_lm_head_test.py#L25
def test_dtype(dtype):
  head = keras_nlp.layers.MaskedLMHead(
      vocabulary_size=100,
      activation="softmax",
      dtype=dtype, # this is the point
  )
  encoded_tokens = tf.keras.Input(shape=(10, 16))
  positions = tf.keras.Input(shape=(5,), dtype="int32")
  outputs = head(encoded_tokens, mask_positions=positions)
  model = tf.keras.Model((encoded_tokens, positions), outputs)

  token_data = tf.random.uniform(shape=(4, 10, 16))
  position_data = tf.random.uniform(minval=0, maxval=10, shape=(4, 5), dtype=tf.int32)
  model((token_data, position_data))

  for w in head.weights:
      assert w.dtype == dtype, ("Wrong type: " + w.name)
      # When it fails it fails with:
      # TypeError: Input 'y' of 'AddV2' Op has type float16 that does not match type float32 of argument 'x'.

print("float32")
test_dtype(tf.float32) # this works

print("bfloat16")
test_dtype(tf.bfloat16) # this fails

print("float64")
test_dtype(tf.float64)

TypeError: Exception encountered when calling layer "masked_lm_head_1" (type MaskedLMHead).

in user code:

    File "/usr/local/lib/python3.10/dist-packages/keras_nlp/src/layers/modeling/masked_lm_head.py", line 196, in call  *
        outputs = outputs + self._bias

    TypeError: Input 'y' of 'AddV2' Op has type bfloat16 that does not match type float32 of argument 'x'.


Call arguments received by layer "masked_lm_head_1" (type MaskedLMHead):
  • inputs=tf.Tensor(shape=(None, 10, 16), dtype=bfloat16)
  • mask_positions=tf.Tensor(shape=(None, 5), dtype=int32)