Support for more complex associative scans

I am trying to implement time-parallel Kalman filters, which can be formulated as a prefix-sum [1]. But the assocative operation is a bit more complicated, so op(::Array, ::Array) is not really sufficient. Most people in the field currently use jax and its associative_scan function, which operates on "pytrees"; typically a tuple of arrays. The scan is then run along a specified axis on these arrays.

Here is a simple example on how associative scans can be used in Jax:

import jax
import jax.numpy as jnp

def double_matmul(elem1: tuple, elem2: tuple):
    Ai, Bi = elem1
    Aj, Bj = elem2
    return Ai @ Aj, Bi @ Bj

def get_elem(dummy_arg=None, D=2):
    # the dummy arg is just there to easily work with vmap
    A = 0.9 * jnp.eye(D)
    B = 1.1 * jnp.eye(D)
    return A, B

# The input for the associative scan will be a tuple of arrays:
As, Bs = jax.vmap(get_elem)(jnp.ones((100)))
# As.shape == Bs.shape == (10,2,2)

Aout, Bout = jax.lax.associative_scan(double_matmul, (As, Bs))
# Aout.shape == Bout.shape == (10,2,2)

Is there a good way to replicate this functionality in Julia with CUDA.jl? Or, is there a possibility to extend CUDA.scan! to enable more complicated inputs, such as tuples of Arrays?

Related: #1482 contains an attempt for running scan! on a CuArray of structs, as well as some discussion on the topic.

[1] Temporal Parallelization of Bayesian Smoothers, Simo Särkkä, Angel F. Garcia-Fernandez (https://arxiv.org/abs/1905.13002)