def batched_swd(X_s, X_t, w_s, w_t, num_chunks=10):
  # Note that this is quick and dirty and only works if the number of samples is divisible by num_chunks
  Xs2 = X_s.reshape(num_chunks, -1, X_s.shape[1])
  ws2 = w_s.reshape(num_chunks, -1)
  Xt2 = X_t.reshape(num_chunks, -1, X_t.shape[1])
  wt2 = w_t.reshape(num_chunks, -1)
  bswdfunc = torch.func.vmap(ot.sliced.sliced_wasserstein_distance,randomness="same")
  swds = bswdfunc(Xs2, Xt2, ws2, wt2)
  return swds.sum()

def approx_swd(
    X_s,
    X_t,
    a=None,
    b=None,
    n_projections=50,
    p=2,
    projections=None,
    seed=None,
    log=False,
    num_bins=100,
):

    X_s, X_t = ot.utils.list_to_array(X_s, X_t)

    if a is not None and b is not None and projections is None:
        nx = ot.utils.get_backend(X_s, X_t, a, b)
    elif a is not None and b is not None and projections is not None:
        nx = ot.utils.get_backend(X_s, X_t, a, b, projections)
    elif a is None and b is None and projections is not None:
        nx = ot.utils.get_backend(X_s, X_t, projections)
    else:
        nx = ot.utils.get_backend(X_s, X_t)

    n = X_s.shape[0]
    m = X_t.shape[0]

    if X_s.shape[1] != X_t.shape[1]:
        raise ValueError(
            "X_s and X_t must have the same number of dimensions {} and {} respectively given".format(
                X_s.shape[1], X_t.shape[1]
            )
        )

    if a is None:
        a = nx.full(n, 1 / n, type_as=X_s)
    if b is None:
        b = nx.full(m, 1 / m, type_as=X_s)

    d = X_s.shape[1]

    if projections is None:
        projections = ot.sliced.get_random_projections(
            d, n_projections, seed, backend=nx, type_as=X_s
        )
    else:
        n_projections = projections.shape[1]

    X_s_projections = nx.dot(X_s, projections)
    X_t_projections = nx.dot(X_t, projections)

    # projected_emd = ot.lp.wasserstein_1d(X_s_projections, X_t_projections, a, b, p=p)

    # Compute histograms
    pminv, pmaxv = torch.aminmax(torch.cat((X_s_projections,X_t_projections)))
    boundaries = torch.linspace(start=pminv, end=pmaxv+torch.finfo(torch.float32).eps, steps=num_bins+1, device=X_s.device)
    bin_idx_s = torch.bucketize(X_s_projections, boundaries[1:])
    bin_idx_t = torch.bucketize(X_t_projections, boundaries[1:])

    centres = 0.5*( boundaries[:-1] + boundaries[1:] )[:,None].expand(-1,n_projections)

    a_rep = a[:,None].expand(-1,n_projections)
    X_s_p_aggr = torch.zeros(num_bins,n_projections, device=X_s.device).scatter_reduce(0, bin_idx_s, a_rep*X_s_projections, reduce="sum", include_self=False)
    a_aggr = torch.zeros(num_bins,n_projections, device=X_s.device).scatter_reduce(0, bin_idx_s, a_rep, reduce="sum", include_self=False)
    X_s_p_aggr /= a_aggr
    zidx = (a_aggr<=torch.finfo(torch.float32).eps)
    X_s_p_aggr[zidx] = centres[zidx]

    b_rep = b[:,None].expand(-1,n_projections)
    X_t_p_aggr = torch.zeros(num_bins,n_projections, device=X_s.device).scatter_reduce(0, bin_idx_t, b_rep*X_t_projections, reduce="sum", include_self=False)
    b_aggr = torch.zeros(num_bins,n_projections, device=X_s.device).scatter_reduce(0, bin_idx_t, b_rep, reduce="sum", include_self=False)
    X_t_p_aggr /= b_aggr
    zidx = (b_aggr<=torch.finfo(torch.float32).eps)
    X_t_p_aggr[zidx] = centres[zidx]

    # Compute EMD on histograms
    projected_emd = ot.lp.wasserstein_1d(X_s_p_aggr, X_t_p_aggr, a_aggr, b_aggr, p=p, require_sort=False)

    res = (nx.sum(projected_emd) / n_projections) ** (1.0 / p)
    if log:
        return res, {"projections": projections, "projected_emds": projected_emd}
    return res