Repository: PennyLaneAI/pennylane

Issue:
[BUG] `AmplitudeEmbedding` with `pad_with` does not work with JAX-JIT
### Expected behavior

It should execute without errors when using `AmplitudeEmbedding` with the `pad_with` argument with JAX-JIT.

### Actual behavior

Code errors out. Jitting works when `pad_with` is `None`.

### Additional information

_No response_

### Source code

```shell
@qml.qnode(dev, interface='jax')
def circ(x):
    qml.AmplitudeEmbedding(x, wires=dev.wires, pad_with=0, normalize=True)
    return qml.state()

circ_jit = jax.jit(circ)

n = np.array([1, 2, 3])
# Works
print(circ(n))
# Does not work
print(circ_jit(n))
```


### Tracebacks

```shell
File "/temp.py", line 25, in <module>
    print(circ_jit(n))
  File "/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/lib/python3.9/site-packages/jax/_src/api.py", line 473, in cache_miss
    out_flat = xla.xla_call(
  File "/lib/python3.9/site-packages/jax/core.py", line 1765, in bind
    return call_bind(self, fun, *args, **params)
  File "/lib/python3.9/site-packages/jax/core.py", line 1781, in call_bind
    outs = top_trace.process_call(primitive, fun_, tracers, params)
  File "/lib/python3.9/site-packages/jax/core.py", line 678, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/lib/python3.9/site-packages/jax/_src/dispatch.py", line 182, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File "/lib/python3.9/site-packages/jax/linear_util.py", line 285, in memoized_fun
    ans = call(fun, *args)
  File "/lib/python3.9/site-packages/jax/_src/dispatch.py", line 230, in _xla_callable_uncached
    return lower_xla_callable(fun, device, backend, name, donated_invars, False,
  File "/lib/python3.9/site-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/lib/python3.9/site-packages/jax/_src/dispatch.py", line 272, in lower_xla_callable
    jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(
  File "/lib/python3.9/site-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 1893, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
  File "/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 1865, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers_)
  File "/lib/python3.9/site-packages/jax/linear_util.py", line 168, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/lib/python3.9/site-packages/pennylane/qnode.py", line 610, in __call__
    self.construct(args, kwargs)
  File "/lib/python3.9/site-packages/pennylane/qnode.py", line 525, in construct
    self._qfunc_output = self.func(*args, **kwargs)
  File "/temp.py", line 15, in circ
    qml.AmplitudeEmbedding(x, wires=dev.wires, pad_with=0, normalize=True)
  File "/lib/python3.9/site-packages/pennylane/templates/embeddings/amplitude.py", line 129, in __init__
    features = self._preprocess(features, wires, pad_with, normalize)
  File "/lib/python3.9/site-packages/pennylane/templates/embeddings/amplitude.py", line 211, in _preprocess
    feature_set = qml.math.concatenate([feature_set, padding], axis=0)
  File "/lib/python3.9/site-packages/pennylane/math/multi_dispatch.py", line 207, in wrapper
    return fn(*args, **kwargs)
  File "/lib/python3.9/site-packages/pennylane/math/multi_dispatch.py", line 300, in concatenate
    return np.concatenate(values, axis=axis, like=like)
  File "/lib/python3.9/site-packages/autoray/autoray.py", line 85, in do
    return get_lib_fn(backend, fn)(*args, **kwargs)
  File "/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 1637, in concatenate
    _stackable(*arrays) or _check_arraylike("concatenate", *arrays)
  File "/lib/python3.9/site-packages/jax/_src/numpy/util.py", line 298, in _check_arraylike
    raise TypeError(msg.format(fun_name, type(arg), pos))
TypeError: concatenate requires ndarray or scalar arguments, got <class 'list'> at position 1.
```


### System information

```shell
Name: PennyLane
Version: 0.26.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /lib/python3.9/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, retworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning, PennyLane-Lightning-GPU, PennyLane-qiskit, pennylane-qulacs
Platform info:           Linux-5.15.0-50-generic-x86_64-with-glibc2.35
Python version:          3.9.12
Numpy version:           1.22.3
Scipy version:           1.8.0
Installed devices:
- lightning.qubit (PennyLane-Lightning-0.26.0)
- default.gaussian (PennyLane-0.26.0)
- default.mixed (PennyLane-0.26.0)
- default.qubit (PennyLane-0.26.0)
- default.qubit.autograd (PennyLane-0.26.0)
- default.qubit.jax (PennyLane-0.26.0)
- default.qubit.tf (PennyLane-0.26.0)
- default.qubit.torch (PennyLane-0.26.0)
- default.qutrit (PennyLane-0.26.0)
- qiskit.aer (PennyLane-qiskit-0.23.0)
- qiskit.basicaer (PennyLane-qiskit-0.23.0)
- qiskit.ibmq (PennyLane-qiskit-0.23.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.23.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.23.0)
- lightning.gpu (PennyLane-Lightning-GPU-0.23.0)
- qulacs.simulator (pennylane-qulacs-0.16.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:

Resolving patch:
diff --git a/pennylane/templates/embeddings/amplitude.py b/pennylane/templates/embeddings/amplitude.py
--- a/pennylane/templates/embeddings/amplitude.py
+++ b/pennylane/templates/embeddings/amplitude.py
@@ -218,7 +218,7 @@ def _preprocess(features, wires, pad_with, normalize):
                         dat_type = type(feature_set)
                         padding = dat_type(padding).to(feature_set.device)
 
-                    feature_set = qml.math.concatenate([feature_set, padding], axis=0)
+                    feature_set = qml.math.hstack([feature_set, padding])
 
             # normalize
             norm = qml.math.sum(qml.math.abs(feature_set) ** 2)


Test patch:
diff --git a/tests/templates/test_embeddings/test_amplitude.py b/tests/templates/test_embeddings/test_amplitude.py
--- a/tests/templates/test_embeddings/test_amplitude.py
+++ b/tests/templates/test_embeddings/test_amplitude.py
@@ -385,6 +385,39 @@ def test_jax_jit(self, tol, features):
 
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
 
+    @pytest.mark.jax
+    @pytest.mark.parametrize("features", all_features)
+    def test_jax_jit_pad_with(self, tol, features):
+        """Tests jax tensors when using JIT and pad_with."""
+
+        import jax
+        import jax.numpy as jnp
+
+        features = jnp.array(features)
+
+        dev = qml.device("default.qubit", wires=4)
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def circuit_decomposed(features):
+            # need to cast to complex tensor, which is implicitly done in the template
+            state = qml.math.cast(
+                qml.math.hstack([features, qml.math.zeros_like(features)]), np.complex128
+            )
+            qml.QubitStateVector(state, wires=range(4))
+            return qml.state()
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def circuit(x=None):
+            qml.AmplitudeEmbedding(x, list(range(4)), pad_with=0.0)
+            return qml.state()
+
+        res = circuit(features)
+        res2 = circuit_decomposed(features)
+
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
     @pytest.mark.tf
     def test_tf(self, tol):
         """Tests tf tensors."""