@triton.autotune(configs, key=['siz_B', 'T', 'B', 'NH', 'DH'])
@triton.jit
def _forward_sequence_kernel(states_initial, Wx, R, b, states_all,
    gates_all, T: tl.constexpr, NS: tl.constexpr, B: tl.constexpr, NH: tl.
    constexpr, DH: tl.constexpr, NGI: tl.constexpr, NGR: tl.constexpr,
    siz_B: tl.constexpr, OUTPUT_GATES: tl.constexpr, DTYPE: tl.constexpr=tl
    .float32):
    idx_b_NH, idx_b_B = tl.program_id(0), tl.program_id(1)
    str_matWx_NH = T * NGI * B * DH
    str_matWx_T = NGI * B * DH
    str_matStatesAll_NH = (T + 1) * NS * B * DH
    str_matStatesAll_T = NS * B * DH
    str_matGatesAll_NH = T * NGI * B * DH
    str_matGatesAll_T = NGI * B * DH
    matHtrans_initial_ptr = tl.make_block_ptr(base=states_initial + 
        idx_b_NH * NS * B * DH + 0 * B * DH, shape=(B, DH), strides=(DH, 1),
        offsets=(idx_b_B * siz_B, 0), block_shape=(siz_B, DH), order=(0, 1))
    matHtrans = tl.load(matHtrans_initial_ptr).to(tl.float32)
    matCtrans_initial_ptr = tl.make_block_ptr(base=states_initial + 
        idx_b_NH * NS * B * DH + 1 * B * DH, shape=(B, DH), strides=(DH, 1),
        offsets=(idx_b_B * siz_B, 0), block_shape=(siz_B, DH), order=(0, 1))
    matCtrans = tl.load(matCtrans_initial_ptr).to(tl.float32)
    matNtrans_initial_ptr = tl.make_block_ptr(base=states_initial + 
        idx_b_NH * NS * B * DH + 2 * B * DH, shape=(B, DH), strides=(DH, 1),
        offsets=(idx_b_B * siz_B, 0), block_shape=(siz_B, DH), order=(0, 1))
    matNtrans = tl.load(matNtrans_initial_ptr).to(tl.float32)
    matMtrans_initial_ptr = tl.make_block_ptr(base=states_initial + 
        idx_b_NH * NS * B * DH + 3 * B * DH, shape=(B, DH), strides=(DH, 1),
        offsets=(idx_b_B * siz_B, 0), block_shape=(siz_B, DH), order=(0, 1))
    matMtrans = tl.load(matMtrans_initial_ptr).to(tl.float32)
    matHtrans_initial_store_ptr = tl.make_block_ptr(base=states_all + 
        idx_b_NH * str_matStatesAll_NH + 0 * str_matStatesAll_T + 0 * B *
        DH, shape=(B, DH), strides=(DH, 1), offsets=(idx_b_B * siz_B, 0),
        block_shape=(siz_B, DH), order=(0, 1))
    tl.store(matHtrans_initial_store_ptr, matHtrans.to(DTYPE))
    matCtrans_initial_store_ptr = tl.make_block_ptr(base=states_all + 
        idx_b_NH * str_matStatesAll_NH + 0 * str_matStatesAll_T + 1 * B *
        DH, shape=(B, DH), strides=(DH, 1), offsets=(idx_b_B * siz_B, 0),
        block_shape=(siz_B, DH), order=(0, 1))
    tl.store(matCtrans_initial_store_ptr, matCtrans.to(DTYPE))
    matNtrans_initial_store_ptr = tl.make_block_ptr(base=states_all + 
        idx_b_NH * str_matStatesAll_NH + 0 * str_matStatesAll_T + 2 * B *
        DH, shape=(B, DH), strides=(DH, 1), offsets=(idx_b_B * siz_B, 0),
        block_shape=(siz_B, DH), order=(0, 1))
    tl.store(matNtrans_initial_store_ptr, matNtrans.to(DTYPE))
    matMtrans_initial_store_ptr = tl.make_block_ptr(base=states_all + 
        idx_b_NH * str_matStatesAll_NH + 0 * str_matStatesAll_T + 3 * B *
        DH, shape=(B, DH), strides=(DH, 1), offsets=(idx_b_B * siz_B, 0),
        block_shape=(siz_B, DH), order=(0, 1))
    tl.store(matMtrans_initial_store_ptr, matMtrans.to(DTYPE))
    matRtrans_i_ptr = tl.make_block_ptr(base=R + idx_b_NH * DH * NGR * DH +
        0 * DH * DH, shape=(DH, DH), strides=(DH, 1), offsets=(0, 0),
        block_shape=(DH, DH), order=(0, 1))
    matRtrans_i = tl.load(matRtrans_i_ptr)
    matRtrans_f_ptr = tl.make_block_ptr(base=R + idx_b_NH * DH * NGR * DH +
        1 * DH * DH, shape=(DH, DH), strides=(DH, 1), offsets=(0, 0),
        block_shape=(DH, DH), order=(0, 1))
    matRtrans_f = tl.load(matRtrans_f_ptr)
    matRtrans_z_ptr = tl.make_block_ptr(base=R + idx_b_NH * DH * NGR * DH +
        2 * DH * DH, shape=(DH, DH), strides=(DH, 1), offsets=(0, 0),
        block_shape=(DH, DH), order=(0, 1))
    matRtrans_z = tl.load(matRtrans_z_ptr)
    matRtrans_o_ptr = tl.make_block_ptr(base=R + idx_b_NH * DH * NGR * DH +
        3 * DH * DH, shape=(DH, DH), strides=(DH, 1), offsets=(0, 0),
        block_shape=(DH, DH), order=(0, 1))
    matRtrans_o = tl.load(matRtrans_o_ptr)
    vecB_i_ptr = b + idx_b_NH * NGI * DH + 0 * DH + tl.arange(0, DH)
    vecB_i = tl.load(vecB_i_ptr)
    vecB_f_ptr = b + idx_b_NH * NGI * DH + 1 * DH + tl.arange(0, DH)
    vecB_f = tl.load(vecB_f_ptr)
    vecB_z_ptr = b + idx_b_NH * NGI * DH + 2 * DH + tl.arange(0, DH)
    vecB_z = tl.load(vecB_z_ptr)
    vecB_o_ptr = b + idx_b_NH * NGI * DH + 3 * DH + tl.arange(0, DH)
    vecB_o = tl.load(vecB_o_ptr)
    for idx_t in range(T):
        matIxtrans_ptr = tl.make_block_ptr(base=Wx + idx_b_NH *
            str_matWx_NH + idx_t * str_matWx_T + 0 * B * DH, shape=(B, DH),
            strides=(DH, 1), offsets=(idx_b_B * siz_B, 0), block_shape=(
            siz_B, DH), order=(0, 1))
        matIxtrans = tl.load(matIxtrans_ptr)
        matFxtrans_ptr = tl.make_block_ptr(base=Wx + idx_b_NH *
            str_matWx_NH + idx_t * str_matWx_T + 1 * B * DH, shape=(B, DH),
            strides=(DH, 1), offsets=(idx_b_B * siz_B, 0), block_shape=(
            siz_B, DH), order=(0, 1))
        matFxtrans = tl.load(matFxtrans_ptr)
        matZxtrans_ptr = tl.make_block_ptr(base=Wx + idx_b_NH *
            str_matWx_NH + idx_t * str_matWx_T + 2 * B * DH, shape=(B, DH),
            strides=(DH, 1), offsets=(idx_b_B * siz_B, 0), block_shape=(
            siz_B, DH), order=(0, 1))
        matZxtrans = tl.load(matZxtrans_ptr)
        matOxtrans_ptr = tl.make_block_ptr(base=Wx + idx_b_NH *
            str_matWx_NH + idx_t * str_matWx_T + 3 * B * DH, shape=(B, DH),
            strides=(DH, 1), offsets=(idx_b_B * siz_B, 0), block_shape=(
            siz_B, DH), order=(0, 1))
        matOxtrans = tl.load(matOxtrans_ptr)
        matRhtrans_i = tl.dot(matHtrans.to(DTYPE), matRtrans_i)
        matRhtrans_f = tl.dot(matHtrans.to(DTYPE), matRtrans_f)
        matRhtrans_z = tl.dot(matHtrans.to(DTYPE), matRtrans_z)
        matRhtrans_o = tl.dot(matHtrans.to(DTYPE), matRtrans_o)
        matIbar = matIxtrans + matRhtrans_i + vecB_i[None, :]
        matFbar = matFxtrans + matRhtrans_f + vecB_f[None, :]
        matZbar = matZxtrans + matRhtrans_z + vecB_z[None, :]
        matObar = matOxtrans + matRhtrans_o + vecB_o[None, :]
        matLogFplusM = matMtrans + tl.log(tl.sigmoid(matFbar))
        matMtrans_next = tl.where(matNtrans == 0.0, matIbar, tl.maximum(
            matIbar, matLogFplusM))
        matI = tl.exp(matIbar - matMtrans_next)
        matF = tl.exp(matLogFplusM - matMtrans_next)
        matZ = triton_tanh(matZbar)
        matO = tl.sigmoid(matObar)
        matCtrans_next = matF * matCtrans + matI * matZ
        matNtrans_next = tl.maximum(matF * matNtrans + matI, 1.0)
        matHtrans_next = matO * (matCtrans_next / matNtrans_next)
        matHtrans_next_ptr = tl.make_block_ptr(base=states_all + idx_b_NH *
            str_matStatesAll_NH + (idx_t + 1) * str_matStatesAll_T + 0 * B *
            DH, shape=(B, DH), strides=(DH, 1), offsets=(idx_b_B * siz_B, 0
            ), block_shape=(siz_B, DH), order=(0, 1))
        tl.store(matHtrans_next_ptr, matHtrans_next.to(DTYPE))
        matCtrans_next_ptr = tl.make_block_ptr(base=states_all + idx_b_NH *
            str_matStatesAll_NH + (idx_t + 1) * str_matStatesAll_T + 1 * B *
            DH, shape=(B, DH), strides=(DH, 1), offsets=(idx_b_B * siz_B, 0
            ), block_shape=(siz_B, DH), order=(0, 1))
        tl.store(matCtrans_next_ptr, matCtrans_next.to(DTYPE))
        matNtrans_next_ptr = tl.make_block_ptr(base=states_all + idx_b_NH *
            str_matStatesAll_NH + (idx_t + 1) * str_matStatesAll_T + 2 * B *
            DH, shape=(B, DH), strides=(DH, 1), offsets=(idx_b_B * siz_B, 0
            ), block_shape=(siz_B, DH), order=(0, 1))
        tl.store(matNtrans_next_ptr, matNtrans_next.to(DTYPE))
        matMtrans_next_ptr = tl.make_block_ptr(base=states_all + idx_b_NH *
            str_matStatesAll_NH + (idx_t + 1) * str_matStatesAll_T + 3 * B *
            DH, shape=(B, DH), strides=(DH, 1), offsets=(idx_b_B * siz_B, 0
            ), block_shape=(siz_B, DH), order=(0, 1))
        tl.store(matMtrans_next_ptr, matMtrans_next.to(DTYPE))
        if OUTPUT_GATES:
            matGatesItrans_ptr = tl.make_block_ptr(base=gates_all + 
                idx_b_NH * str_matGatesAll_NH + idx_t * str_matGatesAll_T +
                0 * B * DH, shape=(B, DH), strides=(DH, 1), offsets=(
                idx_b_B * siz_B, 0), block_shape=(siz_B, DH), order=(0, 1))
            tl.store(matGatesItrans_ptr, matIbar.to(DTYPE))
            matGatesFtrans_ptr = tl.make_block_ptr(base=gates_all + 
                idx_b_NH * str_matGatesAll_NH + idx_t * str_matGatesAll_T +
                1 * B * DH, shape=(B, DH), strides=(DH, 1), offsets=(
                idx_b_B * siz_B, 0), block_shape=(siz_B, DH), order=(0, 1))
            tl.store(matGatesFtrans_ptr, matFbar.to(DTYPE))
            matGatesZtrans_ptr = tl.make_block_ptr(base=gates_all + 
                idx_b_NH * str_matGatesAll_NH + idx_t * str_matGatesAll_T +
                2 * B * DH, shape=(B, DH), strides=(DH, 1), offsets=(
                idx_b_B * siz_B, 0), block_shape=(siz_B, DH), order=(0, 1))
            tl.store(matGatesZtrans_ptr, matZ.to(DTYPE))
            matGatesOtrans_ptr = tl.make_block_ptr(base=gates_all + 
                idx_b_NH * str_matGatesAll_NH + idx_t * str_matGatesAll_T +
                3 * B * DH, shape=(B, DH), strides=(DH, 1), offsets=(
                idx_b_B * siz_B, 0), block_shape=(siz_B, DH), order=(0, 1))
            tl.store(matGatesOtrans_ptr, matO.to(DTYPE))
        matCtrans = matCtrans_next
        matHtrans = matHtrans_next
        matNtrans = matNtrans_next
        matMtrans = matMtrans_next
