Skip to content

Diverse Metrics implemented in TensorFlow

Diverse Tensorflow customied metrics used in the ChoiceLearn library.

HitRate

Bases: Metric

Compute Hit Rate at k.

Source code in choice_learn/utils/metrics.py 
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
class HitRate(tf.keras.metrics.Metric):
    """Compute Hit Rate at k."""

    def __init__(
        self,
        average_on_trip=False,
        top_k: int = 10,
        name=None,
        axis=-1,
        **kwargs,
    ):
        if name is None:
            name = f"hit_rate_at_{top_k}"
        super().__init__(name=name, **kwargs)
        self.top_k = top_k
        self.hit_rate = self.add_variable(
            shape=(), initializer="zeros", name=f"hit_rate_at_{self.top_k}"
        )
        self.n_evals = self.add_variable(shape=(), initializer="zeros", name="n_evals")
        self.average_on_trip = average_on_trip
        self.axis = axis

    def update_state(self, y_true, y_pred, batch=None):
        """Accumulate statistics for the metric.

        Parameters
        ----------
        y_true : np.ndarray
            Ground Truth value
        y_pred : np.ndarray
            Predicted values
        """
        if self.axis == -1:
            if not len(tf.shape(y_pred)) == 2:
                raise ValueError(f"y_pred must be of shape size 2, is {len(tf.shape(y_pred))}.")
        else:
            y_pred = tf.convert_to_tensor(y_pred)
        y_true = tf.cast(y_true, dtype=tf.int32)

        top_k_indices = tf.math.top_k(y_pred, k=self.top_k).indices  # Shape: (batch_size, top_k)
        hits_per_batch = tf.reduce_any(
            tf.equal(
                tf.cast(top_k_indices, tf.int32),
                tf.cast(tf.expand_dims(y_true, axis=1), tf.int32),
            ),
            axis=1,
        )
        hits = tf.cast(hits_per_batch, tf.float32)
        if batch is not None and self.average_on_trip:
            unique_trips, segment_ids = tf.unique(batch)
            trip_means = tf.math.unsorted_segment_mean(hits, segment_ids, tf.shape(unique_trips)[0])
            self.hit_rate.assign_add(tf.reduce_sum(trip_means))
            self.n_evals.assign_add(tf.cast(tf.shape(unique_trips)[0], self.n_evals.dtype))

        else:
            self.hit_rate.assign_add(tf.reduce_sum(hits))
            self.n_evals.assign_add(tf.cast(tf.shape(y_true)[0], self.n_evals.dtype))

    def result(self):
        """Compute the current metric value.

        Returns
        -------
        float
            Negative Log Likelihood value
        """
        return tf.math.divide_no_nan(self.hit_rate, self.n_evals)

result()

Compute the current metric value.

Returns:

Type Description
float

Negative Log Likelihood value
Source code in choice_learn/utils/metrics.py 
311
312
313
314
315
316
317
318
319
def result(self):
    """Compute the current metric value.

    Returns
    -------
    float
        Negative Log Likelihood value
    """
    return tf.math.divide_no_nan(self.hit_rate, self.n_evals)

update_state(y_true, y_pred, batch=None)

Accumulate statistics for the metric.

Parameters:

Name Type Description Default
y_true ndarray

Ground Truth value

 required
y_pred ndarray

Predicted values

 required
Source code in choice_learn/utils/metrics.py 
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
def update_state(self, y_true, y_pred, batch=None):
    """Accumulate statistics for the metric.

    Parameters
    ----------
    y_true : np.ndarray
        Ground Truth value
    y_pred : np.ndarray
        Predicted values
    """
    if self.axis == -1:
        if not len(tf.shape(y_pred)) == 2:
            raise ValueError(f"y_pred must be of shape size 2, is {len(tf.shape(y_pred))}.")
    else:
        y_pred = tf.convert_to_tensor(y_pred)
    y_true = tf.cast(y_true, dtype=tf.int32)

    top_k_indices = tf.math.top_k(y_pred, k=self.top_k).indices  # Shape: (batch_size, top_k)
    hits_per_batch = tf.reduce_any(
        tf.equal(
            tf.cast(top_k_indices, tf.int32),
            tf.cast(tf.expand_dims(y_true, axis=1), tf.int32),
        ),
        axis=1,
    )
    hits = tf.cast(hits_per_batch, tf.float32)
    if batch is not None and self.average_on_trip:
        unique_trips, segment_ids = tf.unique(batch)
        trip_means = tf.math.unsorted_segment_mean(hits, segment_ids, tf.shape(unique_trips)[0])
        self.hit_rate.assign_add(tf.reduce_sum(trip_means))
        self.n_evals.assign_add(tf.cast(tf.shape(unique_trips)[0], self.n_evals.dtype))

    else:
        self.hit_rate.assign_add(tf.reduce_sum(hits))
        self.n_evals.assign_add(tf.cast(tf.shape(y_true)[0], self.n_evals.dtype))

MRR

Bases: Metric

Compute Mean Reciprocal Rank.

Source code in choice_learn/utils/metrics.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
class MRR(tf.keras.metrics.Metric):
    """Compute Mean Reciprocal Rank."""

    def __init__(
        self,
        average_on_trip=False,
        name="mean_reciprocal_rank",
        axis=-1,
        **kwargs,
    ):
        super().__init__(name=name, **kwargs)
        self.mrr = self.add_variable(shape=(), initializer="zeros", name="mrr")
        self.n_evals = self.add_variable(shape=(), initializer="zeros", name="n_evals")
        self.average_on_trip = average_on_trip
        self.axis = axis

    def update_state(
        self,
        y_true,
        y_pred,
        batch=None,
    ):
        """Accumulate statistics for the metric.

        Parameters
        ----------
        y_true : np.ndarray
            Ground Truth value
        y_pred : np.ndarray
            Predicted values
        """
        if self.axis == -1:
            if not len(tf.shape(y_pred)) == 2:
                raise ValueError(f"y_pred must be of shape size 2, is {len(tf.shape(y_pred))}.")
        else:
            y_pred = tf.convert_to_tensor(y_pred)
        y_true = tf.cast(y_true, dtype=tf.int32)

        ranks = tf.argsort(tf.argsort(-y_pred, axis=1), axis=1) + 1  # Shape: (batch_size, n_items)
        item_batch_indices = tf.stack(
            [tf.range(len(y_true)), y_true], axis=1
        )  # Shape: (batch_size, 2)
        item_ranks = tf.gather_nd(ranks, item_batch_indices)  # Shape: (batch_size,)

        # mean_rank = tf.reduce_sum(tf.cast(1 / item_ranks, dtype=tf.float32), axis=self.axis)
        if batch is not None and self.average_on_trip:
            unique_trips, segment_ids = tf.unique(batch)
            reciprocal_ranks = 1.0 / tf.cast(item_ranks, dtype=tf.float32)
            trip_mrrs = tf.math.unsorted_segment_mean(
                reciprocal_ranks, segment_ids, tf.shape(unique_trips)[0]
            )
            self.mrr.assign_add(tf.reduce_sum(trip_mrrs))
            self.n_evals.assign_add(tf.cast(tf.shape(unique_trips)[0], self.n_evals.dtype))
        else:
            self.mrr.assign(self.mrr + tf.reduce_sum(tf.cast(1 / item_ranks, dtype=tf.float32)))
            self.n_evals.assign_add(tf.cast(tf.shape(y_true)[0], self.n_evals.dtype))

    def result(self):
        """Compute the current metric value.

        Returns
        -------
        float
            Negative Log Likelihood value
        """
        return tf.math.divide_no_nan(self.mrr, self.n_evals)

result()

Compute the current metric value.

Returns:

Type Description
float

Negative Log Likelihood value
Source code in choice_learn/utils/metrics.py 
175
176
177
178
179
180
181
182
183
def result(self):
    """Compute the current metric value.

    Returns
    -------
    float
        Negative Log Likelihood value
    """
    return tf.math.divide_no_nan(self.mrr, self.n_evals)

update_state(y_true, y_pred, batch=None)

Accumulate statistics for the metric.

Parameters:

Name Type Description Default
y_true ndarray

Ground Truth value

 required
y_pred ndarray

Predicted values

 required
Source code in choice_learn/utils/metrics.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def update_state(
    self,
    y_true,
    y_pred,
    batch=None,
):
    """Accumulate statistics for the metric.

    Parameters
    ----------
    y_true : np.ndarray
        Ground Truth value
    y_pred : np.ndarray
        Predicted values
    """
    if self.axis == -1:
        if not len(tf.shape(y_pred)) == 2:
            raise ValueError(f"y_pred must be of shape size 2, is {len(tf.shape(y_pred))}.")
    else:
        y_pred = tf.convert_to_tensor(y_pred)
    y_true = tf.cast(y_true, dtype=tf.int32)

    ranks = tf.argsort(tf.argsort(-y_pred, axis=1), axis=1) + 1  # Shape: (batch_size, n_items)
    item_batch_indices = tf.stack(
        [tf.range(len(y_true)), y_true], axis=1
    )  # Shape: (batch_size, 2)
    item_ranks = tf.gather_nd(ranks, item_batch_indices)  # Shape: (batch_size,)

    # mean_rank = tf.reduce_sum(tf.cast(1 / item_ranks, dtype=tf.float32), axis=self.axis)
    if batch is not None and self.average_on_trip:
        unique_trips, segment_ids = tf.unique(batch)
        reciprocal_ranks = 1.0 / tf.cast(item_ranks, dtype=tf.float32)
        trip_mrrs = tf.math.unsorted_segment_mean(
            reciprocal_ranks, segment_ids, tf.shape(unique_trips)[0]
        )
        self.mrr.assign_add(tf.reduce_sum(trip_mrrs))
        self.n_evals.assign_add(tf.cast(tf.shape(unique_trips)[0], self.n_evals.dtype))
    else:
        self.mrr.assign(self.mrr + tf.reduce_sum(tf.cast(1 / item_ranks, dtype=tf.float32)))
        self.n_evals.assign_add(tf.cast(tf.shape(y_true)[0], self.n_evals.dtype))

MeanRank

Bases: Metric

Compute Mean Rank.

Source code in choice_learn/utils/metrics.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
class MeanRank(tf.keras.metrics.Metric):
    """Compute Mean Rank."""

    def __init__(
        self,
        average_on_trip=False,
        name="mean_rank",
        axis=-1,
        **kwargs,
    ):
        super().__init__(name=name, **kwargs)
        self.mr = self.add_variable(shape=(), initializer="zeros", name="mr")
        self.n_evals = self.add_variable(shape=(), initializer="zeros", name="n_evals")
        self.average_on_trip = average_on_trip
        self.axis = axis

    def update_state(
        self,
        y_true,
        y_pred,
        batch=None,
    ):
        """Accumulate statistics for the metric.

        Parameters
        ----------
        y_true : np.ndarray
            Ground Truth value
        y_pred : np.ndarray
            Predicted values
        """
        if self.axis == -1:
            if not len(tf.shape(y_pred)) == 2:
                raise ValueError(f"y_pred must be of shape size 2, is {len(tf.shape(y_pred))}.")
        else:
            y_pred = tf.convert_to_tensor(y_pred)
        y_true = tf.cast(y_true, dtype=tf.int32)

        ranks = tf.argsort(tf.argsort(-y_pred, axis=1), axis=1) + 1  # Shape: (batch_size, n_items)
        item_batch_indices = tf.stack(
            [tf.range(len(y_true)), y_true], axis=1
        )  # Shape: (batch_size, 2)
        item_ranks = tf.gather_nd(ranks, item_batch_indices)  # Shape: (batch_size,)
        float_rank = tf.cast(item_ranks, dtype=tf.float32)

        if batch is not None and self.average_on_trip:
            unique_trips, segment_ids = tf.unique(batch)
            batch_mr = tf.math.unsorted_segment_mean(
                float_rank, segment_ids, tf.shape(unique_trips)[0]
            )
            self.mr.assign_add(tf.reduce_sum(batch_mr))
            self.n_evals.assign_add(tf.cast(tf.shape(unique_trips)[0], self.n_evals.dtype))
        else:
            self.mr.assign_add(tf.reduce_sum(float_rank))
            self.n_evals.assign_add(tf.cast(tf.shape(y_true)[0], self.n_evals.dtype))

    def result(self):
        """Compute the current metric value.

        Returns
        -------
        float
            Negative Log Likelihood value
        """
        return tf.math.divide_no_nan(self.mr, self.n_evals)

result()

Compute the current metric value.

Returns:

Type Description
float

Negative Log Likelihood value
Source code in choice_learn/utils/metrics.py 
242
243
244
245
246
247
248
249
250
def result(self):
    """Compute the current metric value.

    Returns
    -------
    float
        Negative Log Likelihood value
    """
    return tf.math.divide_no_nan(self.mr, self.n_evals)

update_state(y_true, y_pred, batch=None)

Accumulate statistics for the metric.

Parameters:

Name Type Description Default
y_true ndarray

Ground Truth value

 required
y_pred ndarray

Predicted values

 required
Source code in choice_learn/utils/metrics.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
def update_state(
    self,
    y_true,
    y_pred,
    batch=None,
):
    """Accumulate statistics for the metric.

    Parameters
    ----------
    y_true : np.ndarray
        Ground Truth value
    y_pred : np.ndarray
        Predicted values
    """
    if self.axis == -1:
        if not len(tf.shape(y_pred)) == 2:
            raise ValueError(f"y_pred must be of shape size 2, is {len(tf.shape(y_pred))}.")
    else:
        y_pred = tf.convert_to_tensor(y_pred)
    y_true = tf.cast(y_true, dtype=tf.int32)

    ranks = tf.argsort(tf.argsort(-y_pred, axis=1), axis=1) + 1  # Shape: (batch_size, n_items)
    item_batch_indices = tf.stack(
        [tf.range(len(y_true)), y_true], axis=1
    )  # Shape: (batch_size, 2)
    item_ranks = tf.gather_nd(ranks, item_batch_indices)  # Shape: (batch_size,)
    float_rank = tf.cast(item_ranks, dtype=tf.float32)

    if batch is not None and self.average_on_trip:
        unique_trips, segment_ids = tf.unique(batch)
        batch_mr = tf.math.unsorted_segment_mean(
            float_rank, segment_ids, tf.shape(unique_trips)[0]
        )
        self.mr.assign_add(tf.reduce_sum(batch_mr))
        self.n_evals.assign_add(tf.cast(tf.shape(unique_trips)[0], self.n_evals.dtype))
    else:
        self.mr.assign_add(tf.reduce_sum(float_rank))
        self.n_evals.assign_add(tf.cast(tf.shape(y_true)[0], self.n_evals.dtype))

NegativeLogLikeliHood

Bases: Metric

Compute Negative Loglikelihood.

Parameters:

Name Type Description Default
from_logits bool

Whether provided values are logits or probabilities, by default False

 False
sparse bool

Whether y_true is given as an index or a one-hot, by default False

 False
epsilon float

Lower bound for log(.), by default 1e-10

 1e-10
name str

Name of operation, by default "negative_log_likelihood"

 'negative_log_likelihood'
axis int

axis on which to apply the metric, by default -1

 -1
Source code in choice_learn/utils/metrics.py 
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
class NegativeLogLikeliHood(tf.keras.metrics.Metric):
    """Compute Negative Loglikelihood.

    Parameters
    ----------
    from_logits : bool, optional
        Whether provided values are logits or probabilities, by default False
    sparse : bool, optional
        Whether y_true is given as an index or a one-hot, by default False
    epsilon : float, optional
        Lower bound for log(.), by default 1e-10
    name : str, optional
        Name of operation, by default "negative_log_likelihood"
    axis : int, optional
        axis on which to apply the metric, by default -1
    """

    def __init__(
        self,
        from_logits=False,
        sparse=False,
        average_on_trip=False,
        epsilon=1e-10,
        name="negative_log_likelihood",
        axis=-1,
        **kwargs,
    ):
        """Initialize metric.

        Parameters
        ----------
        from_logits : bool, optional
            Whether provided values are logits or probabilities, by default False
        sparse : bool, optional
            Whether y_true is given as an index or a one-hot, by default False
        epsilon : float, optional
            Lower bound for log(.), by default 1e-10
        average_on_trip: bool, optional
            Whether the metric should be averaged over each trip, by default False
        name : str, optional
            Name of operation, by default "negative_log_likelihood"
        axis : int, optional
            axis on which to apply the metric, by default -1
        """
        super().__init__(name=name, **kwargs)
        self.nll = self.add_variable(shape=(), initializer="zeros", name="neg_ll")
        self.n_evals = self.add_variable(shape=(), initializer="zeros", name="n_evals")
        self.from_logits = from_logits
        self.sparse = sparse
        self.average_on_trip = average_on_trip
        self.epsilon = epsilon
        self.axis = axis

    def update_state(self, y_true, y_pred, batch=None, sample_weight=None):
        """Accumulate statistics for the metric.

        Parameters
        ----------
        y_true : np.ndarray
            Ground Truth value
        y_pred : np.ndarray
            Predicted values
        sample_weight : np.ndarray, optional
            sample wise weight, by default None
        """
        if self.axis == -1:
            if not len(tf.shape(y_pred)) == 2:
                raise ValueError(f"y_pred must be of shape size 2, is {len(tf.shape(y_pred))}.")
        if self.from_logits:  # Apply softmax if utilities are given
            y_pred = tf.nn.softmax(y_pred, axis=self.axis)
        else:
            y_pred = tf.convert_to_tensor(y_pred)
        if self.sparse:  # Create OneHot labels if sparse labels are given
            y_true = tf.one_hot(y_true, depth=tf.shape(y_pred)[self.axis])
        else:
            y_true = tf.cast(y_true, y_pred.dtype)

        # Apply label clipping to avoid log(0) and such issues
        y_pred = tf.clip_by_value(y_pred, self.epsilon, 1.0)
        if sample_weight is None:
            nll_value = -tf.reduce_sum(y_true * tf.math.log(y_pred), axis=self.axis)
        else:
            nll_value = -tf.reduce_sum(
                y_true * tf.math.log(y_pred) * tf.expand_dims(sample_weight, axis=-1),
                axis=self.axis,
            )

        if batch is not None and self.average_on_trip:
            unique_trips, segment_ids = tf.unique(batch)
            trip_nlls = tf.math.unsorted_segment_mean(
                nll_value, segment_ids, tf.shape(unique_trips)[0]
            )
            self.nll.assign_add(tf.reduce_sum(trip_nlls))
            self.n_evals.assign_add(tf.cast(tf.shape(unique_trips)[0], self.n_evals.dtype))
        else:
            self.nll.assign(self.nll + tf.reduce_sum(nll_value))
            if sample_weight is None:
                self.n_evals.assign(self.n_evals + tf.cast(tf.shape(y_true)[0], tf.float32))
            else:
                self.n_evals.assign(self.n_evals + tf.reduce_sum(sample_weight))

    def result(self):
        """Compute the current metric value.

        Returns
        -------
        float
            Negative Log Likelihood value
        """
        return tf.math.divide_no_nan(self.nll, self.n_evals)

__init__(from_logits=False, sparse=False, average_on_trip=False, epsilon=1e-10, name='negative_log_likelihood', axis=-1, **kwargs)

Initialize metric.

Parameters:

Name Type Description Default
from_logits bool

Whether provided values are logits or probabilities, by default False

 False
sparse bool

Whether y_true is given as an index or a one-hot, by default False

 False
epsilon float

Lower bound for log(.), by default 1e-10

 1e-10
average_on_trip

Whether the metric should be averaged over each trip, by default False

 False
name str

Name of operation, by default "negative_log_likelihood"

 'negative_log_likelihood'
axis int

axis on which to apply the metric, by default -1

 -1
Source code in choice_learn/utils/metrics.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def __init__(
    self,
    from_logits=False,
    sparse=False,
    average_on_trip=False,
    epsilon=1e-10,
    name="negative_log_likelihood",
    axis=-1,
    **kwargs,
):
    """Initialize metric.

    Parameters
    ----------
    from_logits : bool, optional
        Whether provided values are logits or probabilities, by default False
    sparse : bool, optional
        Whether y_true is given as an index or a one-hot, by default False
    epsilon : float, optional
        Lower bound for log(.), by default 1e-10
    average_on_trip: bool, optional
        Whether the metric should be averaged over each trip, by default False
    name : str, optional
        Name of operation, by default "negative_log_likelihood"
    axis : int, optional
        axis on which to apply the metric, by default -1
    """
    super().__init__(name=name, **kwargs)
    self.nll = self.add_variable(shape=(), initializer="zeros", name="neg_ll")
    self.n_evals = self.add_variable(shape=(), initializer="zeros", name="n_evals")
    self.from_logits = from_logits
    self.sparse = sparse
    self.average_on_trip = average_on_trip
    self.epsilon = epsilon
    self.axis = axis

result()

Compute the current metric value.

Returns:

Type Description
float

Negative Log Likelihood value
Source code in choice_learn/utils/metrics.py 
107
108
109
110
111
112
113
114
115
def result(self):
    """Compute the current metric value.

    Returns
    -------
    float
        Negative Log Likelihood value
    """
    return tf.math.divide_no_nan(self.nll, self.n_evals)

update_state(y_true, y_pred, batch=None, sample_weight=None)

Accumulate statistics for the metric.

Parameters:

Name Type Description Default
y_true ndarray

Ground Truth value

 required
y_pred ndarray

Predicted values

 required
sample_weight ndarray

sample wise weight, by default None

 None
Source code in choice_learn/utils/metrics.py 
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def update_state(self, y_true, y_pred, batch=None, sample_weight=None):
    """Accumulate statistics for the metric.

    Parameters
    ----------
    y_true : np.ndarray
        Ground Truth value
    y_pred : np.ndarray
        Predicted values
    sample_weight : np.ndarray, optional
        sample wise weight, by default None
    """
    if self.axis == -1:
        if not len(tf.shape(y_pred)) == 2:
            raise ValueError(f"y_pred must be of shape size 2, is {len(tf.shape(y_pred))}.")
    if self.from_logits:  # Apply softmax if utilities are given
        y_pred = tf.nn.softmax(y_pred, axis=self.axis)
    else:
        y_pred = tf.convert_to_tensor(y_pred)
    if self.sparse:  # Create OneHot labels if sparse labels are given
        y_true = tf.one_hot(y_true, depth=tf.shape(y_pred)[self.axis])
    else:
        y_true = tf.cast(y_true, y_pred.dtype)

    # Apply label clipping to avoid log(0) and such issues
    y_pred = tf.clip_by_value(y_pred, self.epsilon, 1.0)
    if sample_weight is None:
        nll_value = -tf.reduce_sum(y_true * tf.math.log(y_pred), axis=self.axis)
    else:
        nll_value = -tf.reduce_sum(
            y_true * tf.math.log(y_pred) * tf.expand_dims(sample_weight, axis=-1),
            axis=self.axis,
        )

    if batch is not None and self.average_on_trip:
        unique_trips, segment_ids = tf.unique(batch)
        trip_nlls = tf.math.unsorted_segment_mean(
            nll_value, segment_ids, tf.shape(unique_trips)[0]
        )
        self.nll.assign_add(tf.reduce_sum(trip_nlls))
        self.n_evals.assign_add(tf.cast(tf.shape(unique_trips)[0], self.n_evals.dtype))
    else:
        self.nll.assign(self.nll + tf.reduce_sum(nll_value))
        if sample_weight is None:
            self.n_evals.assign(self.n_evals + tf.cast(tf.shape(y_true)[0], tf.float32))
        else:
            self.n_evals.assign(self.n_evals + tf.reduce_sum(sample_weight))