def _torch_scaled_dot_product_attention(
        query: np.ndarray, key: np.ndarray, value: np.ndarray, mask: np.ndarray = None
    ) -> np.ndarray:
        """
        Two things:
        1. torch cannot consume np.ndarray, so need to convert to torch.Tensor
        2. torch cpu kernel has no half-precision support, so need to cast to float
        """
        query_torch = torch.tensor(query).to(torch.float32)
        key_torch = torch.tensor(key).to(torch.float32)
        value_torch = torch.tensor(value).to(torch.float32)

        mask_torch = None
        if mask is not None:
            mask_torch = torch.tensor(mask)
            if mask.dtype != bool:
                mask_torch = mask_torch.to(torch.float32)

        return (
            torch.nn.functional.scaled_dot_product_attention(
                query_torch, key_torch, value_torch, mask_torch
            )
            .numpy()
            .astype(query.dtype)
        )

        Two things:
        1. torch cannot consume np.ndarray, so need to convert to torch.Tensor
        2. torch cpu kernel has no half-precision support, so need to cast to float
        
coremltools/converters/mil/mil/ops/tests/iOS18/test_transformers.py