Specify axis of colormap for 3d plot (matplotlib)

I'm using ax.contour3D to plot 3D contour plots in matplotlib. I have a colormap that I want to apply along the X axis, but the default is along the Z axis. I have the following code snippet (problematic code at bottom):

def solve(N, J, mu):
    """Solve u(x,t) using the explicit scheme."""
    # Give initial conditions
    u = np.zeros((N+1, J+1))
    
    # Temporal IC: u(x,0) = 1
    for j in range(1, J):
        u[0][j] = 1

    # Spatial BC: u(0,t) = u(1,t) = 1
    for n in range(0, N):
        u[n][0] = 1
        u[n][-1] = 1
    return u


    for j in range(1, J):
        for n in range(1, N-1):
            u[n+1][j] = u[n][j] + mu*(u[n][j+1] - 2*u[n][j] + u[n][j-1])
    return u

# Constants
Xf = 1
Tf = 1
N = 1000
J = 10
xs = np.linspace(0, Xf, J+1)
ts = np.linspace(0, Tf, N+1)
mu = (ts[1]-ts[0]) / (xs[1] - xs[0])**2  # dt/dx**2
X, T = np.meshgrid(xs, ts)
u = solve(N, J, mu)

# !!! Problematic code !!!
fig = plt.figure(figsize=(4,4), dpi=100)
ax = plt.axes(projection='3d')
ax.contour3D(T, X, u, 50, cmap='coolwarm')
plt.show()

As you can see, the plot colors vertically (along the Z-axis). Is there a way to color from left to right (along the X axis)?