CQCL · gamatos · Mar 20, 2024 · Jan 10, 2024
diff --git a/docs/examples.md b/docs/examples.md
@@ -12,6 +12,10 @@ Below are some use-case notebooks. These both illustrate the flexibility of quja
 - [classification.ipynb](https://github.com/CQCL/qujax/blob/develop/examples/classification.ipynb) - train a quantum circuit for binary classification using data re-uploading.
 - [generative_modelling.ipynb](https://github.com/CQCL/qujax/blob/develop/examples/generative_modelling.ipynb) - uses a parameterised quantum circuit as a generative model for a real life dataset. Trains via stochastic gradient Langevin dynamics on the maximum mean discrepancy between statetensor and dataset.
 
+Experimental (i.e. uses an unstable API which might change in future versions):
+
+- [noise_channel_monte_carlo.ipynb](https://github.com/CQCL/qujax/blob/develop/examples/experimental/noise_channel_monte_carlo.ipynb) - statevector simulation of circuit noise using the Monte-Carlo/quantum trajectories approach.
+
 The [pytket](https://github.com/CQCL/pytket) repository also contains `tk_to_qujax` implementations for some of the above at [pytket-qujax_classification.ipynb](https://github.com/CQCL/pytket/blob/main/examples/pytket-qujax-classification.ipynb), 
 [pytket-qujax_heisenberg_vqe.ipynb](https://github.com/CQCL/pytket/blob/main/examples/pytket-qujax_heisenberg_vqe.ipynb) 
 and [pytket-qujax_qaoa.ipynb](https://github.com/CQCL/pytket/blob/main/examples/pytket-qujax_qaoa.ipynb).
diff --git a/docs/experimental.rst b/docs/experimental.rst
@@ -0,0 +1,8 @@
+Experimental
+=======================
+
+.. toctree::
+    :titlesonly:
+
+    Pure state simulation <experimental/statetensor>
+
diff --git a/docs/experimental/statetensor.rst b/docs/experimental/statetensor.rst
@@ -0,0 +1,13 @@
+Pure state simulation
+=======================
+
+.. toctree::
+    :titlesonly:
+
+    statetensor/get_default_gates
+    statetensor/get_default_operations
+    statetensor/get_params_to_statetensor_func
+    statetensor/get_params
+    statetensor/parse_op
+    statetensor/wrap_parameterised_tensor
+
diff --git a/docs/experimental/statetensor/get_default_gates.rst b/docs/experimental/statetensor/get_default_gates.rst
@@ -0,0 +1,4 @@
+get_default_gates
+==============================================
+
+.. autofunction:: qujax.experimental.statetensor.get_default_gates
diff --git a/docs/experimental/statetensor/get_default_operations.rst b/docs/experimental/statetensor/get_default_operations.rst
@@ -0,0 +1,4 @@
+get_default_operations
+==============================================
+
+.. autofunction:: qujax.experimental.statetensor.get_default_operations
diff --git a/docs/experimental/statetensor/get_params.rst b/docs/experimental/statetensor/get_params.rst
@@ -0,0 +1,4 @@
+get_params
+==============================================
+
+.. autofunction:: qujax.experimental.statetensor.get_params
diff --git a/docs/experimental/statetensor/get_params_to_statetensor_func.rst b/docs/experimental/statetensor/get_params_to_statetensor_func.rst
@@ -0,0 +1,4 @@
+get_params_to_statetensor_func
+=====================================
+
+.. autofunction:: qujax.experimental.statetensor.get_params_to_statetensor_func
diff --git a/docs/experimental/statetensor/parse_op.rst b/docs/experimental/statetensor/parse_op.rst
@@ -0,0 +1,4 @@
+parse_op
+==============================================
+
+.. autofunction:: qujax.experimental.statetensor.parse_op
diff --git a/docs/experimental/statetensor/wrap_parameterised_tensor.rst b/docs/experimental/statetensor/wrap_parameterised_tensor.rst
@@ -0,0 +1,4 @@
+wrap_parameterised_tensor
+==============================================
+
+.. autofunction:: qujax.experimental.statetensor.wrap_parameterised_tensor
diff --git a/docs/index.rst b/docs/index.rst
@@ -59,6 +59,7 @@ Contents
     Pure state simulation <statetensor>
     Mixed state simulation <densitytensor>
     Utility functions <utils>
+    Experimental <experimental>
 
 .. toctree::
     :caption: Links:

diff --git a/docs/statetensor/apply_gate.rst b/docs/statetensor/apply_gate.rst
@@ -1,4 +1,4 @@
 apply_gate
 ==============================================
 
-.. autofunction:: qujax.apply_gate
+.. autofunction:: qujax.experimental.statetensor.apply_gate
diff --git a/examples/experimental/noise_channel_monte_carlo.ipynb b/examples/experimental/noise_channel_monte_carlo.ipynb
diff --git a/qujax/__init__.py b/qujax/__init__.py
@@ -41,6 +41,9 @@
 
 import qujax.typing
 
+import qujax.experimental
+import qujax.experimental.statetensor
+
 # pylint: disable=undefined-variable
 del version
 del statetensor

diff --git a/qujax/densitytensor.py b/qujax/densitytensor.py
@@ -18,7 +18,7 @@
 from qujax.typing import (
     MixedCircuitFunction,
     KrausOp,
-    GateFunction,
+    ParameterizedGateFunction,
     GateParameterIndices,
 )
 
@@ -86,7 +86,7 @@ def kraus(
 def _to_kraus_operator_seq_funcs(
     kraus_op: KrausOp,
     param_inds: Optional[Union[GateParameterIndices, Sequence[GateParameterIndices]]],
-) -> Tuple[Sequence[GateFunction], Sequence[jax.Array]]:
+) -> Tuple[Sequence[ParameterizedGateFunction], Sequence[jax.Array]]:
     """
     Ensures Kraus operators are a sequence of functions that map (possibly empty) parameters to
     tensors and that each element of param_inds_seq is a sequence of arrays that correspond to the

diff --git a/qujax/experimental/__init__.py b/qujax/experimental/__init__.py