Function selective_scan_flop_jit

model_zigma.py:1173–1180

(inputs, outputs)

def selective_scan_flop_jit(inputs, outputs):
    # print_jit_input_names(inputs)
    B, D, L = inputs[0].type().sizes()
    N = inputs[2].type().sizes()[1]
    flops = flops_selective_scan_fn(
        B=B, L=L, D=D, N=N, with_D=True, with_Z=False, with_Group=True
    )
    return flops


def flops(model, shape=(3, 32, 32)):