Allow specifying a branch to switch to when setting up a dev environment from the wheels (#1406)

**Context:** Switching to a different branch after setting up the dev
wheel environment would break the latter.

**Description of the Change:** Accept a branch to switch to as an
optional argument, and switch to it before hard-linking the frontend
source code files.

**Benefits:** Start working directly on an existing branch.

**Possible Drawbacks:** Switching to another branch using git would keep
breaking the installation.

Author: rauletorresc

doc/dev/installation.rst

@@ -481,6 +481,12 @@ If you require the Catalyst repository with all its submodules, clone it this wa
 
   git clone --recurse-submodules --shallow-submodules git@github.com:PennyLaneAI/catalyst.git
 
+If you need to work on an existing branch, provide its name as a second argument:
+
+.. code-block:: console
+
+  bash ./setup_dev_from_wheel.sh /path/to/virtual/env branch-name
+
 How Does it Work?
 ^^^^^^^^^^^^^^^^^
 
@@ -495,8 +501,7 @@ using the installed Catalyst wheel libraries, hence avoiding compilation.
 Further Steps
 ^^^^^^^^^^^^^
 
-If everything goes well, ``git status`` should not report any changed files.
-
+``git status`` should not report any changed files when a branch name is not specified.
 Before making changes to the frontend, make sure you create a new branch:
 
 .. code-block:: console
@@ -505,6 +510,9 @@ Before making changes to the frontend, make sure you create a new branch:
 
 Once in the new branch, make the wanted changes. Use the IDE of your preference.
 
+When specifying a branch to switch to, ``git status`` might report changes in some files.
+This is normal. Proceed to make changes in the selected branch.
+
 You can test the changes by executing your sample code under the same virtual environment you used
 with the scripts. As files in the repository are hard-linked to the Wheel code, you are actually
 changing the code stored at the Python ``site-packages`` folder as well, and you will be automatically

doc/releases/changelog-dev.md

@@ -73,6 +73,9 @@
 
   The algorithm is taken from [Relaxed Peephole Optimization: A Novel Compiler Optimization for Quantum Circuits, by Ji Liu, Luciano Bello, and Huiyang Zhou](https://arxiv.org/abs/2012.07711).
 
+* Allow specifying a branch to switch to when setting up a dev environment from the wheels.
+  [(#1406)](https://github.com/PennyLaneAI/catalyst/pull/1406)
+
 <h3>Breaking changes 💔</h3>
 
 * The `sample` and `counts` measurement primitives now support dynamic shot values across catalyst, although at the PennyLane side, the device shots still is constrained to a static integer literal.

setup_dev_from_wheel.sh

@@ -3,6 +3,9 @@
 # Python environment path
 PYTHON_ENV_PATH=$1
 
+# Branch to switch to
+BRANCH=$2
+
 # Exit on any error
 set -e
 
@@ -69,6 +72,11 @@ checkout_nightly_build(){
 
 link_repo_to_wheel(){
     echo "Linking Catalyst repository to Catalyst Wheel..."
+    
+    # switch to branch if given
+    if [ ! -z "${BRANCH}" ]; then
+        git switch $BRANCH
+    fi
 
     export SITEPKGS=$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')
     export CATALYST_WHEEL=$SITEPKGS/catalyst
@@ -81,7 +89,7 @@ restore_catalyst_config(){
     # After linking the Wheel sources, _configuration.py will contain the entry: 'INSTALLED=True'.
     # Hence, we restore the file from the repository.
     cd $CATALYST_DIR
-    git checkout frontend/catalyst/_configuration.py
+    git checkout frontend/catalyst/_configuration.py frontend/catalyst/_version.py
 }
 
 report_changed_files(){