Mixed precision tutorial does not work for custom training loop

[CaptainDario]

I am trying to follow the mixed precision tutorial when using a custom training loop, specifically

@tf.function
def train_step(x, y):
  with tf.GradientTape() as tape:
    predictions = model(x)
    loss = loss_object(y, predictions)
    scaled_loss = optimizer.get_scaled_loss(loss)
  scaled_gradients = tape.gradient(scaled_loss, model.trainable_variables)
  gradients = optimizer.get_unscaled_gradients(scaled_gradients)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

When trying to invoke, get_scaled_loss and get_unscaled_gradients an exception is thrown that those methods are undefined.

Maybe this is related to #19244