bfs errors on sorting #178

import argparse
import time
import psutil
import torch
import wandb
import platform
from cayleypy import CayleyGraph, PermutationGroups


def run_single_n(
    k: int,
    n: int,
    generator_family: str,
    device: str,
    coset: bool,
    central_mode: str,
    inverse_closed: bool,
    coinc: bool,
):
    wandb.init(
        entity="CayleyPy",
        project="cycles",
        name=(
            f"k_{k}_n_{n}_{generator_family}_{device}_"
            f"coset_{int(coset)}_{central_mode if coset else 'full'}_"
            f"inv_{int(inverse_closed)}_coinc_{int(coinc)}"
        ),
        config={
            "k": k,
            "n": n,
            "generator_family": generator_family,
            "device": device,
            "coset": coset,
            "central_mode": central_mode,
            "inverse_closed": inverse_closed,
            "coinc": coinc,
        },
    )

    print(
        f"Running k={k} n={n} gen={generator_family} device={device} "
        f"coset={coset} central_mode={central_mode} inverse_closed={inverse_closed} coinc={coinc}"
    )

    proc = psutil.Process()
    if device == "cuda":
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA requested but not available")
        torch.cuda.reset_peak_memory_stats()
    mem_before = proc.memory_info().rss

    if device == "cuda":
        print(f"Hardware: GPU = {torch.cuda.get_device_name(0)}")
    else:
        print(f"Hardware: CPU = {platform.processor()}, cores = {psutil.cpu_count(logical=True)}")

    t0 = time.time()

    if generator_family == "consecutive":
        defn = PermutationGroups.consecutive_k_cycles(n, k)
    elif generator_family == "wrapped":
        defn = PermutationGroups.wrapped_k_cycles(n, k)
    else:
        raise ValueError(f"Invalid generator_family: {generator_family}")

    if inverse_closed:
        defn = defn.make_inverse_closed()

    # Best-effort: if your cayleypy exposes this, keep/disable coincidences in the generator set.
    if not coinc and hasattr(defn, "remove_coincidences"):
        defn = defn.remove_coincidences()

    # IMPORTANT: full graph vs coset
    if coset:
        if central_mode == "alternating":
            central = [0, 1] * (n // 2) + [0] * (n - 2 * (n // 2))
        elif central_mode == "block":
            central = [0] * (n // 2) + [1] * (n - n // 2)
        else:
            raise ValueError(f"Invalid central_mode: {central_mode}")
        defn = defn.with_central_state(central)

    graph = CayleyGraph(defn)
    result = graph.bfs(return_all_edges=False, return_all_hashes=False)

    last_layer = result.last_layer()
    diameter = result.diameter()
    layer_sizes = result.layer_sizes
    runtime = time.time() - t0

    last_layer_list = [list(row) for row in last_layer]
    last_layer_str = "\n".join(" ".join(map(str, row)) for row in last_layer_list)

    if device == "cuda":
        torch.cuda.synchronize()
        peak_bytes = torch.cuda.max_memory_allocated()
    else:
        mem_after = proc.memory_info().rss
        peak_bytes = max(mem_before, mem_after)

    print(f"n={n}, diameter: {diameter}, layer sizes: {layer_sizes}")
    print(f"Last layer:\n{last_layer_str}")
    print(f"Peak memory: {peak_bytes / 1024**3:.3f} GiB")
    print(f"Runtime: {runtime:.3f} seconds")

    wandb.log(
        dict(
            diameter=diameter,
            num_layers=len(layer_sizes),
            layer_sizes=layer_sizes,
            runtime_sec=runtime,
            peak_memory_bytes=peak_bytes,
            peak_memory_gib=peak_bytes / 1024**3,
            last_layer_str=last_layer_str,
            last_layer_list=last_layer_list,
            k=k,
            n=n,
            generator_family=generator_family,
            device=device,
            coset=coset,
            central_mode=central_mode if coset else None,
            inverse_closed=inverse_closed,
            coinc=coinc,
        )
    )
    wandb.finish()


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--k", type=int, required=True)
    p.add_argument("--n", type=int, required=True)
    p.add_argument("--device", choices=["cpu", "cuda"], default="cuda")
    p.add_argument("--generator_family", choices=["consecutive", "wrapped"], default="consecutive")
    p.add_argument("--coset", action="store_true", help="restrict to a coset via a central state")
    p.add_argument("--central_mode", choices=["alternating", "block"], default="block")
    p.add_argument("--inverse_closed", action="store_true")
    p.add_argument("--coinc", action="store_true", help="allow coincident generators (if supported)")
    return p.parse_args()


if __name__ == "__main__":
    a = parse_args()
    run_single_n(
        a.k,
        a.n,
        a.generator_family,
        a.device,
        a.coset,
        a.central_mode,
        a.inverse_closed,
        a.coinc,
    )

Running k=6 n=14 gen=consecutive device=cuda coset=False central_mode=block inverse_closed=False coinc=False
Hardware: GPU = NVIDIA H200
Traceback (most recent call last):
  File "/gpfs/projects/mathai/vilin/cayleypy/cayleypy/experiments/consecutive_k_cycles.py", line 144, in <module>
    run_single_n(
    ~~~~~~~~~~~~^
        a.k,
        ^^^^
    ...<6 lines>...
        a.coinc,
        ^^^^^^^^
    )
    ^
  File "/gpfs/projects/mathai/vilin/cayleypy/cayleypy/experiments/consecutive_k_cycles.py", line 84, in run_single_n
    result = graph.bfs(return_all_edges=False, return_all_hashes=False)
  File "/gpfs/projects/mathai/vilin/miniconda3/envs/cayley/lib/python3.13/site-packages/cayleypy/cayley_graph.py", line 294, in bfs
    layer2_hashes, _ = torch.sort(layer2_hashes)
                       ~~~~~~~~~~^^^^^^^^^^^^^^^
RuntimeError: The dimension being sorted can not have more than INT_MAX elements.