@triton.jit
def _naive_associative_rnn_scan_fwd_kernel(x_ptr, a_ptr, out_ptr,
    stride_x_batch, stride_x_len, stride_x_dim, stride_a_batch,
    stride_a_len, stride_a_dim, stride_out_batch, stride_out_len,
    stride_out_dim, seq_len: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    pid_batch = tl.program_id(0)
    pid_dim = tl.program_id(1)
    x_ptr += pid_batch * stride_x_batch
    a_ptr += pid_batch * stride_a_batch
    out_ptr += pid_batch * stride_out_batch
    offsets_dim = pid_dim * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    offsets_len = tl.arange(0, seq_len)
    x_ptrs = x_ptr + offsets_dim[None, :] * stride_x_dim + offsets_len[:, None
        ] * stride_x_len
    a_ptrs = a_ptr + offsets_dim[None, :] * stride_a_dim + offsets_len[:, None
        ] * stride_a_len
    out_ptrs = out_ptr + offsets_dim[None, :] * stride_out_dim + offsets_len[
        :, None] * stride_out_len
    x = tl.load(x_ptrs).to(tl.float32)
    a = tl.load(a_ptrs).to(tl.float32)
    _, all_hiddens = tl.associative_scan(input=(a, x), axis=0, combine_fn=
        _associative_scan_op)
    tl.store(out_ptrs, all_hiddens.to(out_ptr.dtype.element_ty), mask=(
        offsets_len == seq_len - 1)[:, None])
