def batch_directions(z, y, Q, path_sizes, step_vals, subtract_projection=True):
    """
    This function takes an input batch of z vectors (and corresponding class label vectors y)
    and applies the directions in Q to the batch of z vectors.

    :param z: (N, nz) tensor of base random noise to manipulate with directions
    :param y: (N, D) tensor of class vectors
    :param Q: (ndirs, nz) matrix of z-space directions
    :param path_sizes: (ndirs,) tensor indicating how far to travel in each direction
    :param step_vals: (interp_steps,) tensor controlling the granularity of the interpolation
    :param subtract_projection: bool, whether or not to "remove" each direction from the sampled z vectors
    :return: z: (N * ndirs * interp_steps, nz) tensor, y: (N * ndirs * interp_steps) tensor
             containing all z's and y's needed to create the visualizations
    """
    interp_steps = step_vals.size(0)
    N, nz = z.size()
    ndirs = Q.size(0)
    z = z.view(1, N, 1, nz).repeat(ndirs, 1, interp_steps, 1)  # .view(N * ndirs * interp_steps, nz)
    if subtract_projection:
        # The projection will be the same across the interp_steps dimension, so we can just pick-out the first step:
        z_proj = z[:, :, 0, :].view(ndirs * N, nz)
        Q_proj = Q.repeat_interleave(N, dim=0)
        projection = (z_proj * Q_proj).sum(dim=1, keepdims=True) / Q_proj.pow(2).sum(dim=1, keepdims=True) * Q_proj
        z -= projection.view(ndirs, N, 1, nz)
    path_sizes = path_sizes.view(ndirs, 1, 1, 1)
    step_vals = step_vals.view(1, 1, interp_steps, 1)
    Q = Q.view(ndirs, 1, 1, nz)
    z += step_vals * path_sizes * Q
    z = z.view(N * ndirs * interp_steps, nz)
    y = y.repeat_interleave(interp_steps, dim=0).repeat(ndirs, 1)
    return z, y