def MambaInnerFnNoOutProj_jit(inputs, outputs):
            """
            conv1d_weight = rearrange(conv1d_weight, "d 1 w -> d w")
            x, z = xz.chunk(2, dim=1)
            conv1d_out = causal_conv1d_cuda.causal_conv1d_fwd(x, conv1d_weight, conv1d_bias,None, True)
            x_dbl = F.linear(rearrange(conv1d_out, 'b d l -> (b l) d'), x_proj_weight)  # (bl d)
            delta = rearrange(delta_proj_weight @ x_dbl[:, :delta_rank].t(), "d (b l) -> b d l", l = L)
            B = x_dbl[:, delta_rank:delta_rank + d_state]
            B = rearrange(B, "(b l) dstate -> b 1 dstate l", l=L).contiguous()
            C = x_dbl[:, -d_state:]
            C = rearrange(C, "(b l) dstate -> b 1 dstate l", l=L).contiguous()
            out, scan_intermediates, out_z = selective_scan_cuda.fwd(
                conv1d_out, delta, A, B, C, D, z, delta_bias, delta_softplus
            )
            """
            xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight, A = inputs[0:6]
            Batch, _, L = xz.type().sizes()
            CWidth = conv1d_weight.type().sizes()[-1]
            H = A.type().sizes()[-1] # 16
            Dim, R = delta_proj_weight.type().sizes()
            assert tuple(xz.type().sizes()) == (Batch, 2 * Dim, L)
            assert tuple(conv1d_weight.type().sizes()) == (Dim, 1, CWidth)
            assert tuple(x_proj_weight.type().sizes()) == (R + H + H, Dim)
            assert tuple(A.type().sizes()) == (Dim, H)

            with_Z = True
            with_D = False
            if "D" in inputs[6].debugName():
                assert tuple(inputs[6].type().sizes()) == (Dim,)
                with_D = True

            flops = 0
            flops += Batch * (Dim * L) * CWidth # causal_conv1d_cuda.causal_conv1d_fwd
            flops += Batch * (Dim * L) * (R + H + H) # x_dbl = F.linear(...
            flops += Batch * (Dim * R) * (L) # delta_proj_weight @ x_dbl[:, :delta_rank]
            
            # https://github.com/state-spaces/mamba/issues/110
            flops = 9 * Batch * L * Dim * H
            if with_D:
                flops += Batch * Dim * L
            if with_Z:
                flops += Batch * Dim * L  

            return flops

            conv1d_weight = rearrange(conv1d_weight, "d 1 w -> d w")
            x, z = xz.chunk(2, dim=1)
            conv1d_out = causal_conv1d_cuda.causal_conv1d_fwd(x, conv1d_weight, conv1d_bias,None, True)
            x_dbl = F.linear(rearrange(conv1d_out, 'b d l -> (b l) d'), x_proj_weight)  # (bl d)
            delta = rearrange(delta_proj_weight @ x_dbl[:, :delta_rank].t(), "d (b l) -> b d l", l = L)
            B = x_dbl[:, delta_rank:delta_rank + d_state]
            B = rearrange(B, "(b l) dstate -> b 1 dstate l", l=L).contiguous()
            C = x_dbl[:, -d_state:]
            C = rearrange(C, "(b l) dstate -> b 1 dstate l", l=L).contiguous()
            out, scan_intermediates, out_z = selective_scan_cuda.fwd(
                conv1d_out, delta, A, B, C, D, z, delta_bias, delta_softplus
            )
            
analyze/analyze_for_vim.py