Make jit_dispatch aware of the inner grid length