Function build_api_selective_scan

test_selective_scan_easy.py:844–850

(chunksize=64)

# api to fit original mamba_ssm
def build_api_selective_scan(chunksize=64):
    def selective_scan_fn(u, delta, A, B, C, D, z=None,
        delta_bias=None, delta_softplus=None,
        return_last_state=False):
        assert z is None
        return selective_scan_easy(u, delta, A, B, C, D, delta_bias, delta_softplus, return_last_state, chunksize)
    return selective_scan_fn


def selective_scan_ref(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False,