[BUG] qml.matrix on qutrit circuits is giving the wrong sized matrix

Expected behavior

When calling qml.matrix on a qutrit qnode the resulting matrix should be (3^n x 3^n) (n is the number of wires).

For example when there are two wires the resulting matrix should be 9x9.

Actual behavior

When calling qml.matrix on a qutrit qnode with only single-qubit gates gives a matrix with shape (3(2^(n-1)) x 3(2^(n-1))) (n is the number of wires). It doesn't matter if there are gates on one or more wires.

When calling qml.matrix on a qutrit qnode with only two-qubit gates gives a matrix with shape (9(2^(n-2)) x 9(2^(n-2))) (n is the number of wires).

Calling qml.matrix on a qnode with a two-qutrit and one-qutrit gate is causing an error due to the different sized matrices that are produced.

Additional information

Circuits run correctly.

Source code

>>> dev = qml.device("default.qutrit", wires = 2)
>>> @qml.qnode(device = dev)
>>> def circ():
>>>     qml.THadamard(wires=0)
>>>     return qml.state()
>>> print(circ())
[0.-0.57735027j 0.+0.j         0.+0.j         0.-0.57735027j
 0.+0.j         0.+0.j         0.-0.57735027j 0.+0.j
 0.+0.j        ]
>>> print(np.shape(qml.matrix(circ)()))
(6, 6)

-----------------------------------------------------------------------------------------

>>> dev = qml.device("default.qutrit", wires = 2)
>>> @qml.qnode(device = dev)
>>> def circ():
>>>     qml.THadamard(wires=0)
>>>     qml.THadamard(wires=1)
>>>     return qml.state()
>>> print(circ())
[-0.33333333+0.j -0.33333333+0.j -0.33333333+0.j -0.33333333+0.j
 -0.33333333+0.j -0.33333333+0.j -0.33333333+0.j -0.33333333+0.j
 -0.33333333+0.j]
>>> print(np.shape(qml.matrix(circ)()))
(6, 6)

Tracebacks

No response

System information

Name: PennyLane
Version: 0.32.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /usr/local/lib/python3.10/dist-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning

Platform info:           Linux-5.15.109+-x86_64-with-glibc2.31
Python version:          3.10.12
Numpy version:           1.22.4
Scipy version:           1.10.0
Installed devices:

Existing GitHub issues

[X] I have searched existing GitHub issues to make sure the issue does not already exist.

PennyLaneAI / pennylane