Use of precision at recall loss from Eban et al in Keras

Question

I'd like to test non standard losses such as precision_at_recall_loss described in https://arxiv.org/abs/1608.04802 using keras.

Those losses are implemented in loss_layers.py and util.py here: https://github.com/tensorflow/models/tree/archive/research/global_objectives

The following code is a demo using the MNIST dataset.

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

import loss_layers
import util 

def precision_recall_auc_loss(y_true, y_pred):
    y_true = keras.backend.reshape(y_true, (batch_size, 1)) 
    y_pred = keras.backend.reshape(y_pred, (batch_size, 1))   
    util.get_num_labels = lambda labels : 1
    return loss_layers.precision_recall_auc_loss(y_true, y_pred)[0]

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)
input_shape = x_train.shape[1:]

num_classes = 10
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

model = keras.Sequential([
        layers.Conv2D(32,kernel_size=(3,3),activation='relu',input_shape=input_shape), \
        layers.MaxPooling2D(2,2), \
        layers.Flatten(), \
        layers.Dropout(0.25), \
        layers.Dense(num_classes, activation="softmax")
    ])

model.summary()
    
batch_size = 30
epochs = 10
target_recall = 0.9

model.compile(loss=precision_recall_auc_loss,
            optimizer=keras.optimizers.Adam(lr=0.001))

model.fit(x_train, y_train, batch_size=batch_size, \
          epochs=epochs, validation_split=0.15)

The model compiles and start fitting. However, I get the following error:

Train on 51000 samples, validate on 9000 samples
Epoch 1/10
FailedPreconditionError: Attempting to use uninitialized value precision_at_recall_1/lambdas
     [[{{node precision_at_recall_1/lambdas/read}}]]