Visualizing circuits with mid-circuit measurement/feedforward

[abidart]

I am trying to visualize circuits for a new chapter in Qbook. I have tried with two different approaches: with cirq_utils and with tsim.

def show_circuit_cirq(squin_kernel, *args):
    return emit_circuit(squin_kernel, args=args, ignore_returns=True)

def show_circuit_tsim(squin_kernel, height=400):
    @squin.kernel
    def _to_visualize():
        _ = squin_kernel()

    return tsim.Circuit(_to_visualize).diagram(height=height)

Neither of these supports mid-circuit measurement/feedforward, but when run on a simple teleportation squin kernel, the tsim version fails whereas the cirq_utils version silently skips the mid-circuit measurements/feedforward steps and prints the circuit as if they were not there.

Example teleportation circuit

@squin.kernel
def teleport() -> Register:
    input_register = psi()
    bell_register = bell_state_prep()

    squin.cx(input_register[0], bell_register[0])
    squin.h(input_register[0])

    mid_circuit_measurement = squin.broadcast.measure([input_register[0], bell_register[0]])

    if mid_circuit_measurement[0]:
        squin.z(bell_register[1])
    if mid_circuit_measurement[1]:
        squin.x(bell_register[1])
    
    return [bell_register[1]]

show_circuit_cirq(teleport)
show_circuit_tsim(teleport)

Is there a correct way to visualize squin kernels with mid-circuit measurement/feedforward operations?