Tweak AUC test to avoid sharing state across calls in np.apply_along_…

…axis.

This was causing issues when porting over to the new test decorators. The root
cause is unclear. It's possible that m.reset_states() isn't sufficiently
thorough. In any case, avoiding sharing the state altogether solves the issue.
Since np.apply_along_axis doesn't seem to gain anything by reusing the
`tf.keras.metrics.AUC` instance across calls, it seems safest to avoid reusing it.

PiperOrigin-RevId: 274245423

tensorflow_probability/python/stats/ranking_test.py

@@ -67,13 +67,11 @@ def testAurocAuprc(self, curve, positive_means, negative_means):
     y_true = np.array([1] * num_positive_trials + [0] * num_negative_trials)
     y_pred = np.concatenate([positive_trials_, negative_trials_])
 
-    m = tf.keras.metrics.AUC(num_thresholds=total_trials, curve=curve)
-    self.evaluate([v.initializer for v in m.variables])
-
     def auc_fn(y_pred):
+      m = tf.keras.metrics.AUC(num_thresholds=total_trials, curve=curve)
+      self.evaluate([v.initializer for v in m.variables])
       self.evaluate(m.update_state(y_true, y_pred))
       out = self.evaluate(m.result())
-      m.reset_states()
       return out
 
     batch_shape = np.array(positive_means).shape