Repository: scikit-learn/scikit-learn
File: sklearn/calibration.py (Python)

Code before:
            """
            # TODO: numpy 2.0
            # Ensure raw_prediction has the same dtype as labels using .astype().
            # Without this, dtype promotion rules differ across NumPy versions:
            #
            #   beta = np.float64(0)
            #   logits = np.array([1, 2], dtype=np.float32)
            #
            #   result = beta * logits
            #   - NumPy < 2: result.dtype is float32
            #   - NumPy 2+:  result.dtype is float64
            #
            #  This can cause dtype mismatch errors downstream (e.g., buffer dtype).
            if _is_numpy_namespace(xp):
                raw_prediction = (np.exp(log_beta) * logits).astype(dtype_)
                return halfmulti_loss(
                    y_true=labels,
                    raw_prediction=raw_prediction,
                    sample_weight=sample_weight,
                )

            log_beta = xp.asarray(log_beta, dtype=dtype_, device=xp_device)
            raw_prediction = xp.exp(log_beta) * logits
            return _half_multinomial_loss(
                y=labels, pred=raw_prediction, sample_weight=sample_weight, xp=xp
            )

        xatol = 64 * xp.finfo(dtype_).eps
        log_beta_minimizer = minimize_scalar(
            log_loss,
            bounds=(-10.0, 10.0),
            options={
                "xatol": xatol,
            },
        )

        if not log_beta_minimizer.success:  # pragma: no cover
            raise RuntimeError(
                "Temperature scaling fails to optimize during calibration. "
                "Reason from `scipy.optimize.minimize_scalar`: "
                f"{log_beta_minimizer.message}"
            )

        self.beta_ = xp.exp(
            xp.asarray(log_beta_minimizer.x, dtype=dtype_, device=xp_device)
        )

        return self

    def predict(self, X):
        """Predict new data by linear interpolation.

Review comment (suggestion):
We can probably call `_is_numpy_namespace(xp)` just once, instead of doing so in every iteration of the optimisation process.

We could initialise the half multinomial loss as follows:

```python
if _is_numpy_namespace(xp):
    loss = HalfMultinomialLoss(n_classes=logits.shape[1])

else:
    loss = partial(_half_multinomial_loss, xp=xp)
```
Then `log_loss` would look like this:

```python
def log_loss(log_beta=0.0):
    log_beta = xp.asarray(log_beta, dtype=dtype_, device=xp_device)
    raw_prediction = xp.exp(log_beta) * logits
    return loss(labels, raw_prediction, sample_weight)
```

The (CUDA) CI passed locally on my machine.

Code after:

            Returns
            -------
            negative_log_likelihood_loss : float
                The negative log likelihood loss.

            """
            # TODO: numpy 2.0
            # Ensure raw_prediction has the same dtype as labels using .astype().
            # Without this, dtype promotion rules differ across NumPy versions:
            #
            #   beta = np.float64(0)
            #   logits = np.array([1, 2], dtype=np.float32)
            #
            #   result = beta * logits
            #   - NumPy < 2: result.dtype is float32
            #   - NumPy 2+:  result.dtype is float64
            #
            #  This can cause dtype mismatch errors downstream (e.g., buffer dtype).
            log_beta = xp.asarray(log_beta, dtype=dtype_, device=xp_device)
            raw_prediction = xp.exp(log_beta) * logits
            return multinomial_loss(labels, raw_prediction, sample_weight)

        xatol = 64 * xp.finfo(dtype_).eps
        log_beta_minimizer = minimize_scalar(
            log_loss,
            bounds=(-10.0, 10.0),
            options={
                "xatol": xatol,
            },
        )

        if not log_beta_minimizer.success:  # pragma: no cover
            raise RuntimeError(
                "Temperature scaling fails to optimize during calibration. "
                "Reason from `scipy.optimize.minimize_scalar`: "
                f"{log_beta_minimizer.message}"
            )

        self.beta_ = xp.exp(
            xp.asarray(log_beta_minimizer.x, dtype=dtype_, device=xp_device)
        )

        return self

    def predict(self, X):
        """Predict new data by linear interpolation.

        Parameters
        ----------
        X : ndarray of shape (n_samples,) or (n_samples, n_classes)