import numpy as np
import torch

from sklearn.metrics import zero_one_loss
from sklearn.base import config_context

with config_context(array_api_dispatch=True):
    print(zero_one_loss(
        torch.tensor([[0, 1], [1, 1], [1, 0]], dtype=torch.float64),
        torch.tensor(np.ones((3, 2)), dtype=torch.float64),
        sample_weight=torch.tensor([0, 3, 4], dtype=torch.float64)
    ))