Dx not working with 1D mixed function space about shenfun HOT 4 OPEN

from sympy import symbols, pi
from sympy.functions import cos, sin
from shenfun import inner, div, grad, TestFunction, TrialFunction, Function, \
    ETDRK4, FunctionSpace, Array, \
    comm, get_simplified_tpmatrices, Dx, MixedFunctionSpace, la

# Use sympy to set up initial condition
x = symbols("x", real=True)
u_1_0 = cos(x*pi)
u_2_0 = sin(x*pi)

# Initial conditions
u0 = 20*cos(x*pi)
v0 = 20*sin(x*pi)

# Parameters
a = 2
b = 6
D = 10

# Size of discretization
T1 = FunctionSpace(100, 'F', dtype='D', domain=(0, 20))
T2 = FunctionSpace(100, 'F', dtype='D', domain=(0, 20))
T = MixedFunctionSpace([T1, T2])
u = TrialFunction(T)
v = TestFunction(T)
padding_factor = 1.5

# Declare solution arrays and work arrays
UV = Array(T, buffer=(u_1_0, u_2_0))
UV_hat = Function(T)
w0 = Function(T)       # Work array spectral space
Tp = T.get_dealiased(padding_factor=1.5)
w1 = Array(Tp)

#initialize
UV_hat = UV.forward(UV_hat)

u0, u1 = u
v0, v1 = v
dt = 0.01
A = inner(v0, u0 - dt*Dx(u0, 0, 2)) + inner(v1, u1-dt*Dx(u1, 0, 2))
sol = la.BlockMatrixSolver(A)

def NonlinearRHS(uv_hat, rhs):
    rhs.fill(0)
    UVp = uv_hat.backward(padding_factor=padding_factor) # 3/2-rule dealiasing for nonlinear term
    w1[0] = a + UVp[0]**2*UVp[1] - (b+1)*UVp[0]
    w1[1] = b*UVp[0] - UVp[0]**2*UVp[1]
    rhs += w1.forward(rhs)
    return rhs

if __name__ == '__main__':
    rhs = Function(T)
    Nt = 4
    for n in range(Nt):
        rhs[:] = NonlinearRHS(UV_hat, rhs)
        UV_hat = sol(UV_hat + dt*rhs, u=UV_hat)

A = inner(v0, u0 - dt*Dx(u0, 0, 2)) + inner(v1, u1-dt*Dx(u1, 0, 2))
sol = la.BlockMatrixSolver(A)

UV_hat = sol(dt*rhs, u=UV_hat)

UV_hat = sol(UV_hat + dt*rhs, u=UV_hat)