expand_dims eager optimization

doc/dev_start_guide.rst

@@ -168,15 +168,9 @@ create a virtual environment in the project directory:
 
 .. code-block:: bash
 
-    conda env create -n pytensor-dev -f environment.yml
+    conda env create -f environment.yml
     conda activate pytensor-dev
 
-Afterward, you can install the development dependencies:
-
-.. code-block:: bash
-
-    pip install -r requirements.txt
-
 Next, ``pre-commit`` needs to be configured so that the linting and code quality
 checks are performed before each commit:
 
@@ -202,12 +196,12 @@ see the `NumPy development guide <https://numpy.org/doc/stable/dev/>`_.
 Contributing to the documentation
 ---------------------------------
 
-To contribute to the documentation, first follow the instructions in the previous section. Afterward, you can install the documentation dependencies in the virtual environment you created:
-
+The documentation build dependencies have also been included in the virtual environment you created. You can also create a separate virtual environment just for the documentation using the `environment.yml` file located inside the `doc` folder.
 
 .. code-block:: bash
 
-    pip install -r requirements-rtd.txt
+    conda env create -f doc/environment.yml
+    conda activate pytensor-docs
 
 
 You can now build the documentation from the root of the project with:

environment.yml

@@ -37,6 +37,9 @@ dependencies:
   - pygments
   - pydot
   - ipython
+  - pip
+  - pip:
+    - git+https://github.com/pymc-devs/pymc-sphinx-theme
   # code style
   - ruff
   # developer tools

pytensor/tensor/basic.py

@@ -4247,6 +4247,18 @@ def expand_dims(
 
     Insert a new axis that will appear at the `axis` position in the expanded
     array shape.
+
+    Parameters
+    ----------
+    a :
+        The input array.
+    axis :
+        Position in the expanded axes where the new axis is placed.
+        If `axis` is ``None``, `a` will be returned immediately.
+
+    Returns
+    -------
+    `a` with a new axis at the `axis` position.
     """
     a = as_tensor(a)
 
@@ -4256,6 +4268,9 @@ def expand_dims(
     out_ndim = len(axis) + a.ndim
     axis = np.core.numeric.normalize_axis_tuple(axis, out_ndim)
 
+    if not axis:
+        return a
+
     dim_it = iter(range(a.ndim))
     pattern = ["x" if ax in axis else next(dim_it) for ax in range(out_ndim)]