def _tsne(
    X: Union[torch.Tensor, np.ndarray],
    no_dims: int = 2,
    initial_dims: int = 50,
    perplexity: float = 30.0,
    max_iter: int = 1000,
    verbose: bool = False,
    device=None,
) -> Tensor:
    """
    Runs t-SNE on the dataset in the NxD array X to reduce its
    dimensionality to no_dims dimensions. The syntaxis of the function is
    `Y = tsne.tsne(X, no_dims, perplexity), where X is an NxD NumPy array.
    """

    if not isinstance(no_dims, int) or no_dims <= 0:
        raise ValueError("dims must be positive integer")
    if isinstance(X, np.ndarray):
        X = torch.tensor(X, device=device)

    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if verbose:
        print(f"using {device}", file=sys.stderr)

    X.to(device)

    if verbose:
        print("initializing...", file=sys.stderr)
    # Initialize variables
    if initial_dims < X.shape[1]:
        X = _pca_torch(X, device, initial_dims, verbose=verbose)
    else:
        if verbose:
            print(
                "skipping PCA because initial_dims is larger than input dimensionality",
                file=sys.stderr,
            )
    (n, d) = X.shape
    initial_momentum = 0.5
    final_momentum = 0.8
    eta = 500
    min_gain = 0.01
    Y = torch.randn(n, no_dims, device=device)
    iY = torch.zeros(n, no_dims, device=device)
    gains = torch.ones(n, no_dims, device=device)

    # Compute P-values
    if verbose:
        print("computing p-values...", file=sys.stderr)
    P = _x2p_torch(X, 1e-5, initial_dims, perplexity, device=device, verbose=verbose)
    P = P + P.t()
    P = P / (torch.sum(P) + 1e-10)
    P = P * 4.0  # early exaggeration
    P = torch.max(P, torch.tensor(1e-21, device=device))  # (N, N)

    if verbose:
        print("fitting...", file=sys.stderr)

    bar = range(max_iter)

    if verbose:
        bar = tqdm(bar)

    for it in bar:
        # Compute pairwise affinities
        sum_Y = torch.sum(Y * Y, dim=1)
        num = -2.0 * torch.mm(Y, Y.t())  # (N, N)
        num = 1.0 / (1.0 + (num + sum_Y).t() + sum_Y)
        num.fill_diagonal_(0)
        Q = num / (torch.max(torch.sum(num), torch.tensor(1e-12, device=device)))
        Q = torch.max(Q, torch.tensor(1e-12, device=device))

        # Compute gradient
        PQ = P - Q
        # ((N, 1, N).repeat -> (N, no_dims, N).permute -> (N, N, no_dims)) *
        # ((N, no_dims) - (N, 1, no_dims) -> (N, N, no_dims))
        # -> (N, N, no_dims).sum[1] -> (N, no_dims)
        dY = torch.sum(
            (PQ * num).unsqueeze(1).repeat(1, no_dims, 1).transpose(2, 1)
            * (Y.unsqueeze(1) - Y),
            dim=1,
        )

        # Perform the update
        if it < 20:
            momentum = initial_momentum
        else:
            momentum = final_momentum

        gains = (gains + 0.2) * ((dY > 0.0) != (iY > 0.0)).float() + (gains * 0.8) * (
            (dY > 0.0) == (iY > 0.0)
        ).float()
        gains[gains < min_gain] = min_gain
        iY = momentum * iY - eta * (gains * dY)
        Y = Y + iY
        Y = Y - torch.mean(Y, 0)

        # Compute current value of cost function
        if verbose:
            C = torch.sum(P * torch.log(P / Q))
            bar.set_description(f"error: {C.cpu().item():.3f}")
            # if C.item() <= 0.8:
            #     break

        # Stop lying about P-values
        if it == 100:
            P = P / 4.0

    # Return solution
    return Y

    Runs t-SNE on the dataset in the NxD array X to reduce its
    dimensionality to no_dims dimensions. The syntaxis of the function is
    `Y = tsne.tsne(X, no_dims, perplexity), where X is an NxD NumPy array.
    
smoe/utils/visualization/tsne_torch_model.py