penalize_func

    def equal_f(self, inp, val):
        res = tf.reduce_sum(tf.cast(tf.equal(inp, val), tf.int32)) + 1
        return res

    def call(self, inp, training=None, y=None):
        seq_length = tf.math.reduce_sum(tf.cast(tf.math.not_equal(inp, 0), dtype=tf.int32), axis=1)
        mask = tf.cast(tf.math.not_equal(inp, 0), tf.float32)[:, tf.newaxis, tf.newaxis, :]  #(batch_size, 1, 1, seq_len)
        enc_output = self.encoder(inp, training, mask)  # (batch_size, inp_seq_len, d_model)
        logits = self.final_layer(enc_output)  # (batch_size, tar_seq_len, target_vocab_size)

        if y is not None:
            range_labels = tf.range(0, self.target_vocab_size)
            res = lambda val: self.equal_f(y, val)
            count_arr = tf.cast(tf.map_fn(res, range_labels), dtype=tf.float32)
            priores = count_arr / tf.reduce_sum(count_arr)
            logits = tf.nn.log_softmax(logits) - tf.math.log(priores)

        pred = tf.cast(tf.argmax(logits, axis=-1), dtype=tf.int32)
        return pred, logits, seq_length