Issues with `static_argnums` and `static_argnames` when using a pytree argument

Seems to be a issues with qjit's static_argnames and static_argnums when an argument is a multi-argument pytree:

@qml.qjit(static_argnames="y")
def f(x, y):
    if y:
        return x['y']
    return x['x']

@qml.qjit(static_argnums=1)
def g(x, y):
    if y:
        return x['y']
    return x['x']

>>> f({"x": 0.0, "y": 1.0}, y="truthy")
CompileError: argnum 1 is beyond the valid range of [0, 1).```
>>> f({"x": 0.0, "y": 1.0}, "truthy")
TypeError: Argument 'truthy' of type <class 'str'> is not a valid JAX type
>>> g({"x": 0.0, "y": 1.0}, y="truthy")
CompileError: argnum 1 is beyond the valid range of [0, 1).

Start by running the reproducer for qjit with static_argnames and static_argnums exactly as shown, then trace how pytree arguments are counted and mapped to static arguments. Compare the reported CompileError and TypeError cases to the intended handling, and add coverage showing that both options work with the demonstrated pytree input.