pytensor 2.27.0__tar.gz → 2.27.1__tar.gz

{pytensor-2.27.0/pytensor.egg-info → pytensor-2.27.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pytensor
-Version: 2.27.0
+Version: 2.27.1
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:

{pytensor-2.27.0 → pytensor-2.27.1}/doc/core_development_guide.rst

@@ -26,4 +26,4 @@ some of them might be outdated though:
 
 * :ref:`unittest` -- Tutorial on how to use unittest in testing PyTensor.
 
-* :ref:`sparse` -- Description of the ``sparse`` type in PyTensor.
+* :ref:`libdoc_sparse` -- Description of the ``sparse`` type in PyTensor.

{pytensor-2.27.0 → pytensor-2.27.1}/doc/extending/creating_a_c_op.rst

@@ -923,7 +923,7 @@ pre-defined macros.  These section tags have no macros: ``init_code``,
 discussed below.
 
 *      ``APPLY_SPECIFIC(str)`` which will automatically append a name
-       unique to the :ref:`Apply` node that applies the `Op` at the end
+       unique to the :ref:`apply` node that applies the `Op` at the end
        of the provided ``str``. The use of this macro is discussed
        further below.
 
@@ -994,7 +994,7 @@ Apply node in their own names to avoid conflicts between the different
 versions of the apply-specific code. The code that wasn't
 apply-specific was simply defined in the ``c_support_code`` method.
 
-To make indentifiers that include the :ref:`Apply` node name use the
+To make indentifiers that include the :ref:`apply` node name use the
 ``APPLY_SPECIFIC(str)`` macro. In the above example, this macro is
 used when defining the functions ``vector_elemwise_mult`` and
 ``vector_times_vector`` as well as when calling function

{pytensor-2.27.0 → pytensor-2.27.1}/doc/extending/creating_an_op.rst

@@ -7,7 +7,7 @@ Creating a new :class:`Op`: Python implementation
 So suppose you have looked through the library documentation and you don't see
 a function that does what you want.
 
-If you can implement something in terms of an existing :ref:`Op`, you should do that.
+If you can implement something in terms of an existing :ref:`op`, you should do that.
 Odds are your function that uses existing PyTensor expressions is short,
 has no bugs, and potentially profits from rewrites that have already been
 implemented.

{pytensor-2.27.0 → pytensor-2.27.1}/doc/extending/inplace.rst

@@ -200,7 +200,7 @@ input(s)'s memory). From there, go to the previous section.
    certainly lead to erroneous computations.
 
    You can often identify an incorrect `Op.view_map` or :attr:`Op.destroy_map`
-   by using :ref:`DebugMode`.
+   by using :ref:`DebugMode <debugmode>`.
 
 .. note::
    Consider using :class:`DebugMode` when developing

{pytensor-2.27.0 → pytensor-2.27.1}/doc/extending/other_ops.rst

@@ -197,7 +197,7 @@ Want C speed without writing C code for your new Op? You can use Numba
 to generate the C code for you! Here is an `example
 Op <https://gist.github.com/nouiz/5492778#file-theano_op-py>`_ doing that.
 
-.. _alternate_PyTensor_types:
+.. _alternate_pytensor_types:
 
 Alternate PyTensor Types
 ========================

{pytensor-2.27.0 → pytensor-2.27.1}/doc/install.rst

@@ -10,7 +10,7 @@ The latest release of PyTensor can be installed from Pypi using `pip`:
     pip install pytensor
 
 
-Or via conda-force:
+Or via conda-forge:
 
 .. code-block:: bash
 

{pytensor-2.27.0 → pytensor-2.27.1}/doc/library/compile/mode.rst

@@ -20,6 +20,9 @@ PyTensor defines the following modes by name:
 
 - ``'FAST_COMPILE'``: Apply just a few graph rewrites and only use Python implementations.
 - ``'FAST_RUN'``: Apply all rewrites, and use C implementations where possible.
+- ``NUMBA``: Apply all relevant related rewrites and compile the whole graph using Numba.
+- ``JAX``: Apply all relevant rewrites and compile the whole graph using JAX.
+- ``PYTORCH`` Apply all relevant rewrites and compile the whole graph using PyTorch compile.
 - ``'DebugMode'``: A mode for debugging. See :ref:`DebugMode <debugmode>` for details.
 - ``'NanGuardMode``: :ref:`Nan detector <nanguardmode>`
 - ``'DEBUG_MODE'``: Deprecated. Use the string DebugMode.
@@ -28,6 +31,12 @@ The default mode is typically ``FAST_RUN``, but it can be controlled via the
 configuration variable :attr:`config.mode`, which can be
 overridden by passing the keyword argument to :func:`pytensor.function`.
 
+For Numba, JAX, and PyTorch, we exclude rewrites that introduce C-only Ops,
+as well as BLAS optimizations, as those are done automatically by the respective backends.
+
+For JAX we also exclude fusion and inplace optimizations, as JAX does not support them
+at the user level. They are performed automatically by JAX.
+
 .. TODO::
 
     For a finer level of control over which rewrites are applied, and whether

{pytensor-2.27.0 → pytensor-2.27.1}/doc/library/tensor/random/index.rst

@@ -83,7 +83,7 @@ Low-level objects
 .. automodule:: pytensor.tensor.random.op
    :members: RandomVariable, default_rng
 
-..automodule:: pytensor.tensor.random.type
+.. automodule:: pytensor.tensor.random.type
    :members: RandomType, RandomGeneratorType, random_generator_type
 
 .. automodule:: pytensor.tensor.random.var

{pytensor-2.27.0 → pytensor-2.27.1}/doc/troubleshooting.rst

@@ -145,44 +145,64 @@ How do I configure/test my BLAS library
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 There are many ways to configure BLAS for PyTensor. This is done with the PyTensor
-flags ``blas__ldflags`` (:ref:`libdoc_config`). The default is to use the BLAS
-installation information in NumPy, accessible via
-``numpy.__config__.show()``.  You can tell pytensor to use a different
-version of BLAS, in case you did not compile NumPy with a fast BLAS or if NumPy
-was compiled with a static library of BLAS (the latter is not supported in
-PyTensor).
+flags ``blas__ldflags`` (:ref:`libdoc_config`). If not specified, PyTensor will
+attempt to find a local BLAS library to link against, prioritizing specialized implementations.
+The details can be found in :func:`pytensor.link.c.cmodule.default_blas_ldflags`.
 
-The short way to configure the PyTensor flags ``blas__ldflags`` is by setting the
-environment variable :envvar:`PYTENSOR_FLAGS` to ``blas__ldflags=XXX`` (in bash
-``export PYTENSOR_FLAGS=blas__ldflags=XXX``)
+Users can manually set the PyTensor flags ``blas__ldflags`` to link against a
+specific version. This is useful even if the default version is the desired one,
+as it will avoid the costly work of trying to find the best BLAS library at runtime.
 
-The ``${HOME}/.pytensorrc`` file is the simplest way to set a relatively
-permanent option like this one.  Add a ``[blas]`` section with an ``ldflags``
-entry like this:
+The PyTensor flags can be set in a few ways:
+
+1. In the ``${HOME}/.pytensorrc`` file.
 
 .. code-block:: cfg
 
     # other stuff can go here
     [blas]
-    ldflags = -lf77blas -latlas -lgfortran #put your flags here
+    ldflags = -llapack -lblas -lcblas  # put your flags here
 
     # other stuff can go here
 
-For more information on the formatting of ``~/.pytensorrc`` and the
-configuration options that you can put there, see :ref:`libdoc_config`.
+2. In BASH before running your script:
+
+.. code-block:: bash
+
+    export PYTENSOR_FLAGS="blas__ldflags='-llapack -lblas -lcblas'"
+
+3. In an Ipython/Jupyter notebook before importing PyTensor:
+
+.. code-block:: python
+
+    %set_env PYTENSOR_FLAGS=blas__ldflags='-llapack -lblas -lcblas'
+
+
+4. In `pytensor.config` directly:
+
+.. code-block:: python
+
+    import pytensor
+    pytensor.config.blas__ldflags = '-llapack -lblas -lcblas'
+
+
+(For more information on the formatting of ``~/.pytensorrc`` and the
+configuration options that you can put there, see :ref:`libdoc_config`.)
+
+You can find the default BLAS library that PyTensor is linking against by
+checking ``pytensor.config.blas__ldflags``
+or running :func:`pytensor.link.c.cmodule.default_blas_ldflags`.
 
 Here are some different way to configure BLAS:
 
-0) Do nothing and use the default config, which is to link against the same
-BLAS against which NumPy was built. This does not work in the case NumPy was
-compiled with a static library (e.g. ATLAS is compiled by default only as a
-static library).
+0) Do nothing and use the default config.
+This will usually work great for installation via conda/mamba/pixi (conda-forge channel).
+It will usually fail to link altogether for installation via pip.
 
 1) Disable the usage of BLAS and fall back on NumPy for dot products. To do
-this, set the value of ``blas__ldflags`` as the empty string (ex: ``export
-PYTENSOR_FLAGS=blas__ldflags=``). Depending on the kind of matrix operations your
-PyTensor code performs, this might slow some things down (vs. linking with BLAS
-directly).
+this, set the value of ``blas__ldflags`` as the empty string.
+Depending on the kind of matrix operations your PyTensor code performs,
+this might slow some things down (vs. linking with BLAS directly).
 
 2) You can install the default (reference) version of BLAS if the NumPy version
 (against which PyTensor links) does not work. If you have root or sudo access in
@@ -208,10 +228,29 @@ correctly (for example, for MKL this might be ``-lmkl -lguide -lpthread`` or
 ``-lmkl_intel_lp64 -lmkl_intel_thread -lmkl_core -lguide -liomp5 -lmkl_mc
 -lpthread``).
 
+5) Use another backend such as Numba or JAX that perform their own BLAS optimizations,
+by setting the configuration mode to ``"NUMBA"`` or ``"JAX"`` and making sure those packages are installed.
+This configuration mode can be set in all the ways that the BLAS flags can be set, described above.
+
+Alternatively, you can pass `mode='NUMBA'` when compiling individual PyTensor functions without changing the default.
+or use the ``config.change_flags`` context manager.
+
+.. code-block:: python
+
+    from pytensor import function, config
+    from pytensor.tensor import matrix
+
+    x = matrix('x')
+    y = x @ x.T
+    f = function([x], y, mode='NUMBA')
+
+    with config.change_flags(mode='NUMBA'):
+        # compiling function that benefits from BLAS using NUMBA
+        f = function([x], y)
+
 .. note::
 
-    Make sure your BLAS
-    libraries are available as dynamically-loadable libraries.
+    Make sure your BLAS libraries are available as dynamically-loadable libraries.
     ATLAS is often installed only as a static library.  PyTensor is not able to
     use this static library. Your ATLAS installation might need to be modified
     to provide dynamically loadable libraries.  (On Linux this
@@ -267,7 +306,7 @@ configuration information. Then, it will print the running time of the same
 benchmarks for your installation. Try to find a CPU similar to yours in
 the table, and check that the single-threaded timings are roughly the same.
 
-PyTensor should link to a parallel version of Blas and use all cores
+PyTensor should link to a parallel version of BLAS and use all cores
 when possible. By default it should use all cores. Set the environment
 variable "OMP_NUM_THREADS=N" to specify to use N threads.
 

{pytensor-2.27.0 → pytensor-2.27.1}/doc/tutorial/examples.rst

@@ -347,15 +347,7 @@ afterwards compile this expression to get functions,
 using pseudo-random numbers is not as straightforward as it is in
 NumPy, though also not too complicated.
 
-The way to think about putting randomness into PyTensor's computations is
-to put random variables in your graph. PyTensor will allocate a NumPy
-`RandomStream` object (a random number generator) for each such
-variable, and draw from it as necessary. We will call this sort of
-sequence of random numbers a *random stream*. *Random streams* are at
-their core shared variables, so the observations on shared variables
-hold here as well. PyTensor's random objects are defined and implemented in
-:ref:`RandomStream<libdoc_tensor_random_utils>` and, at a lower level,
-in :ref:`RandomVariable<libdoc_tensor_random_basic>`.
+The general user-facing API is documented in :ref:`RandomStream<libdoc_tensor_random_basic>`
 
 For a more technical explanation of how PyTensor implements random variables see :ref:`prng`.
 

{pytensor-2.27.0 → pytensor-2.27.1}/doc/tutorial/prng.rst

@@ -31,8 +31,8 @@ In the first line np.random.default_rng(seed) creates a random Generator.
 >>> rng # doctest: +SKIP
 Generator(PCG64) at 0x7F6C04535820
 
-Every numpy Generator holds a BitGenerator, which is able to generate high-quality sequences of pseudo random bits.
-Numpy generators convert these sequences of bits into sequences of numbers that follow a specific statistical distribution.
+Every NumPy Generator holds a BitGenerator, which is able to generate high-quality sequences of pseudo random bits.
+NumPy generators' methods convert these sequences of bits into sequences of numbers that follow a specific statistical distribution.
 For more details, you can read `NumPy random sampling documentation <https://numpy.org/doc/stable/reference/random>`_.
 
 >>> rng.bit_generator # doctest: +SKIP
@@ -47,6 +47,7 @@ For more details, you can read `NumPy random sampling documentation <https://num
 
 When we call rng.uniform(size=2), the Generator class requested a new array of pseudo random bits (state) from the BitGenerator,
 and used a deterministic mapping function to convert those into a float64 numbers.
+
 It did this twice, because we requested two draws via the size argument.
 In the long-run this deterministic mapping function should produce draws that are statistically indistinguishable from a true uniform distribution.
 
@@ -71,7 +72,7 @@ array([0.033, 0.972, 0.459, 0.71 , 0.765])
 SciPy
 -----
 
-Scipy wraps these Numpy routines in a slightly different API.
+SciPy wraps these NumPy routines in a slightly different API.
 
 >>> import scipy.stats as st
 >>> rng = np.random.default_rng(seed=123)
@@ -82,7 +83,7 @@ PyTensor
 --------
 
 PyTensor does not implement its own bit/generators methods.
-Just like Scipy, it borrows NumPy routines directly.
+Just like SciPy, it borrows NumPy routines directly.
 
 The low-level API of PyTensor RNGs is similar to that of SciPy,
 whereas the higher-level API of RandomStreams is more like that of NumPy.
@@ -95,20 +96,19 @@ We will look at RandomStreams shortly, but we will start with the low-level API.
 >>> x = pt.random.uniform(size=2, rng=rng)
 >>> f = pytensor.function([rng], x)
 
-We created a function that takes a Numpy RandomGenerator and returns two uniform draws. Let's evaluate it
+We created a function that takes a NumPy RandomGenerator and returns two uniform draws. Let's evaluate it
 
 >>> rng_val = np.random.default_rng(123)
 >>> print(f(rng_val), f(rng_val))
 [0.68235186 0.05382102] [0.68235186 0.05382102]
 
-The first numbers were exactly the same as the numpy and scipy calls, because we are using the very same routines.
+The first numbers were exactly the same as the NumPy and SciPy calls, because we are using the very same routines.
 
 Perhaps surprisingly, we got the same results when we called the function the second time!
 This is because PyTensor functions do not hold an internal state and do not modify inputs inplace unless requested to.
 
-We made sure that the rng_val was not modified when calling our Pytensor function, by copying it before using it.
-This may feel inefficient (and it is), but PyTensor is built on a pure functional approach, which is not allowed to have side-effects
-(such as changing global variables) by default.
+We made sure that the rng_val was not modified when calling our PyTensor function, by copying it before using it.
+This may feel inefficient (and it is), but PyTensor is built on a pure functional approach, which is not allowed to have side-effects by default.
 
 We will later see how we can get around this issue by making the inputs mutable or using shared variables with explicit update rules.
 
@@ -129,8 +129,8 @@ In this case we had to advance it twice to get two completely new draws, because
 But other distributions could need more states for a single draw, or they could be clever and reuse the same state for multiple draws.
 
 Because it is not in generally possible to know how much one should modify the generator's bit generator,
-PyTensor RandomVariables actually return the copied generator as a hidden output.
-This copied generator can be safely used again because it contains the bit generator that was already modified when taking draws.
+PyTensor RandomVariables actually return the used generator as a hidden output.
+This generator can be safely used again because it contains the bit generator that was already modified when taking draws.
 
 >>> next_rng, x = x.owner.outputs
 >>> next_rng.type, x.type
@@ -148,7 +148,6 @@ uniform_rv{"(),()->()"}.0 [id A] <RandomGeneratorType> 'next_rng'
     └─ 1.0 [id G] <Scalar(float32, shape=())>
 uniform_rv{"(),()->()"}.1 [id A] <Vector(float64, shape=(2,))> 'x'
  └─ ···
-<ipykernel.iostream.OutStream at 0x7fa5d3a475e0>
 
 We can see the single node with [id A], has two outputs, which we named next_rng and x. By default only the second output x is given to the user directly, and the other is "hidden".
 
@@ -226,14 +225,14 @@ This is exactly what RandomStream does behind the scenes
 >>> x.owner.inputs[0], x.owner.inputs[0].default_update  # doctest: +SKIP
 (RNG(<Generator(PCG64) at 0x7FA45F4A3760>), uniform_rv{"(),()->()"}.0)
 
-From the example here, you can see that RandomStream uses a NumPy-like API in contrast to
-the SciPy-like API of `pytensor.tensor.random`. Full documentation can be found at
-:doc:`../library/tensor/random/basic`.
-
 >>> f = pytensor.function([], x)
 >>> print(f(), f(), f())
 0.19365083425294516 0.7541389670292019 0.2762903411491048
 
+From the example here, you can see that RandomStream uses a NumPy-like API in contrast to
+the SciPy-like API of `pytensor.tensor.random`. Full documentation can be found at
+:doc:`libdoc_tensor_random_basic`.
+
 Shared RNGs are created by default
 ----------------------------------
 
@@ -279,7 +278,7 @@ RandomStreams provide a helper method to achieve the same
 Inplace optimization
 ====================
 
-As mentioned before, by default RandomVariables return a copy of the next RNG state, which can be quite slow.
+As mentioned, RandomVariable Ops default to making a copy of the input RNG before using it, which can be quite slow.
 
 >>> rng = np.random.default_rng(123)
 >>> rng_shared = pytensor.shared(rng, name="rng")
@@ -291,13 +290,13 @@ uniform_rv{"(),()->()"}.1 [id A] 'x' 0
  ├─ NoneConst{None} [id C]
  ├─ 0.0 [id D]
  └─ 1.0 [id E]
-<ipykernel.iostream.OutStream at 0x7fa5d3a475e0>
+
 
 >>> %timeit f()  # doctest: +SKIP
-169 µs ± 24.6 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
+81.8 µs ± 15.4 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
 
 >>> %timeit rng.uniform()  # doctest: +SKIP
-3.56 µs ± 106 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
+2.15 µs ± 63.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
 
 Like other PyTensor operators, RandomVariable's can be given permission to modify inputs inplace during their operation.
 
@@ -307,16 +306,6 @@ If the flag is set, the RNG will not be copied before taking random draws.
 >>> x.owner.op.inplace
 False
 
-This flag is printed as the last argument of the Op in the `dprint`
-
->>> pytensor.dprint(x) # doctest: +SKIP
-uniform_rv{"(),()->()"}.1 [id A] 'x' 0
- ├─ rng [id B]
- ├─ NoneConst{None} [id C]
- ├─ 0.0 [id D]
- └─ 1.0 [id E]
-<ipykernel.iostream.OutStream at 0x7fa5d3a475e0>
-
 For illustration purposes, we will subclass the Uniform Op class and set inplace to True by default.
 
 Users should never do this directly!
@@ -336,27 +325,21 @@ uniform_rv{"(),()->()"}.1 [id A] d={0: [0]} 0
  ├─ NoneConst{None} [id C]
  ├─ 0.0 [id D]
  └─ 1.0 [id E]
-<ipykernel.iostream.OutStream at 0x7fa5d3a475e0>
 
-The destroy map annotation tells us that the first output of the x variable is allowed to alter the first input.
+The destroy map annotation tells us that the first output of the x variable is allowed to modify the first input.
 
 >>> %timeit inplace_f() # doctest: +SKIP
-35.5 µs ± 1.87 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
+9.71 µs ± 2.06 µs per loop (mean ± std. dev. of 7 runs, 100000 loops each)
 
-Performance is now much closer to calling numpy directly, with only a small overhead introduced by the PyTensor function.
+Performance is now much closer to calling NumPy directly, with a small overhead introduced by the PyTensor function.
 
 The `random_make_inplace <https://github.com/pymc-devs/pytensor/blob/3fcf6369d013c597a9c964b2400a3c5e20aa8dce/pytensor/tensor/random/rewriting/basic.py#L42-L52>`_
 rewrite automatically replaces RandomVariable Ops by their inplace counterparts, when such operation is deemed safe. This happens when:
 
 #. An input RNG is flagged as `mutable` and is used in not used anywhere else.
-#. A RNG is created intermediately and used in not used anywhere else.
+#. A RNG is created intermediately and not used anywhere else.
 
-The first case is true when a users uses the `mutable` `kwarg` directly, or much more commonly,
-when a shared RNG is used and a (default or manual) update expression is given.
-In this case, a RandomVariable is allowed to modify the RNG because the shared variable holding it will be rewritten anyway.
-
-The second case is not very common, because RNGs are not usually chained across multiple RandomVariable Ops.
-See more details in the next section.
+The first case is true when a users uses the `mutable` `kwarg` directly.
 
 >>> from pytensor.compile.io import In
 >>> rng = pt.random.type.RandomGeneratorType()("rng")
@@ -371,7 +354,9 @@ uniform_rv{"(),()->()"}.1 [id A] d={0: [0]} 0
  ├─ NoneConst{None} [id C]
  ├─ 0.0 [id D]
  └─ 1.0 [id E]
-<ipykernel.iostream.OutStream at 0x7fa5d3a475e0>
+
+Or, much more commonly, when a shared RNG is used and a (default or manual) update expression is given.
+In this case, a RandomVariable is allowed to modify the RNG because the shared variable holding it will be rewritten anyway.
 
 >>> rng = pytensor.shared(np.random.default_rng(), name="rng")
 >>> next_rng, x = pt.random.uniform(rng=rng).owner.outputs
@@ -385,7 +370,9 @@ uniform_rv{"(),()->()"}.1 [id A] d={0: [0]} 0
  └─ 1.0 [id E]
 uniform_rv{"(),()->()"}.0 [id A] d={0: [0]} 0
  └─ ···
-<ipykernel.iostream.OutStream at 0x7fa5d3a475e0>
+
+The second case is not very common, because RNGs are not usually chained across multiple RandomVariable Ops.
+See more details in the next section.
 
 Multiple random variables
 =========================
@@ -420,7 +407,6 @@ normal_rv{"(),()->()"}.0 [id A] <RandomGeneratorType> 'next_rng_x' 0
  └─ ···
 normal_rv{"(),()->()"}.0 [id F] <RandomGeneratorType> 'next_rng_y' 1
  └─ ···
-<ipykernel.iostream.OutStream at 0x7fa5d3a475e0>
 
 >>> f(), f(), f()
 ([array(-9.8912135), array(-9.80160951)],
@@ -450,7 +436,6 @@ normal_rv{"(),()->()"}.0 [id A] <RandomGeneratorType> 0
  └─ ···
 normal_rv{"(),()->()"}.0 [id F] <RandomGeneratorType> 1
  └─ ···
-<ipykernel.iostream.OutStream at 0x7fa5d3a475e0>
 
 >>> f(), f(), f()
 ([array(-5.81223492), array(-5.85081162)],
@@ -460,15 +445,15 @@ normal_rv{"(),()->()"}.0 [id F] <RandomGeneratorType> 1
 We could have used a single rng.
 
 >>> rng_x = pytensor.shared(np.random.default_rng(seed=123), name="rng_x")
->>> next_rng_x, x = pt.random.normal(loc=0, scale=1, rng=rng).owner.outputs
+>>> next_rng_x, x = pt.random.normal(loc=0, scale=1, rng=rng_x).owner.outputs
 >>> next_rng_x.name = "next_rng_x"
 >>> next_rng_y, y = pt.random.normal(loc=100, scale=1, rng=next_rng_x).owner.outputs
 >>> next_rng_y.name = "next_rng_y"
 >>>
->>> f = pytensor.function([], [x, y], updates={rng: next_rng_y})
+>>> f = pytensor.function([], [x, y], updates={rng_x: next_rng_y})
 >>> pytensor.dprint(f, print_type=True) # doctest: +SKIP
 normal_rv{"(),()->()"}.1 [id A] <Scalar(float64, shape=())> 0
- ├─ rng [id B] <RandomGeneratorType>
+ ├─ rng_x [id B] <RandomGeneratorType>
  ├─ NoneConst{None} [id C] <NoneTypeT>
  ├─ 0 [id D] <Scalar(int8, shape=())>
  └─ 1 [id E] <Scalar(int8, shape=())>
@@ -480,24 +465,23 @@ normal_rv{"(),()->()"}.1 [id F] <Scalar(float64, shape=())> 1
  └─ 1 [id E] <Scalar(int8, shape=())>
 normal_rv{"(),()->()"}.0 [id F] <RandomGeneratorType> 'next_rng_y' 1
  └─ ···
-<ipykernel.iostream.OutStream at 0x7fa5d3a475e0>
 
 >>> f(), f()
-([array(0.91110389), array(101.4795275)],
- [array(0.0908175), array(100.59639646)])
+([array(-0.98912135), array(99.63221335)],
+ [array(1.28792526), array(100.19397442)])
 
-It works, but that graph is slightly unorthodox in Pytensor.
+It works, but that graph is slightly unorthodox in PyTensor.
 
-One practical reason is that it is more difficult to define the correct update expression for the shared RNG variable.
+One practical reason why, is that it is more difficult to define the correct update expression for the shared RNG variable.
 
-One techincal reason is that it makes rewrites more challenging in cases where RandomVariables could otherwise be manipulated independently.
+One techincal reason why, is that it makes rewrites more challenging in cases where RandomVariables could otherwise be manipulated independently.
 
 Creating multiple RNG variables
 -------------------------------
 
 RandomStreams generate high quality seeds for multiple variables, following the NumPy best practices https://numpy.org/doc/stable/reference/random/parallel.html#parallel-random-number-generation.
 
-Users who create their own RNGs should follow the same practice!
+Users who sidestep RandomStreams, either by creating their own RNGs or relying on RandomVariable's default shared RNGs, should follow the same practice!
 
 Random variables in inner graphs
 ================================
@@ -629,7 +613,7 @@ RNGs in Scan are only supported via shared variables in non-sequences at the mom
 >>>     print(err)
 Tensor type field must be a TensorType; found <class 'pytensor.tensor.random.type.RandomGeneratorType'>.
 
-In the future, TensorTypes may be allowed as explicit recurring states, rendering the use of updates optional or unnecessary
+In the future, RandomGenerator variables may be allowed as explicit recurring states, rendering the internal use of updates optional or unnecessary
 
 OpFromGraph
 -----------
@@ -671,7 +655,7 @@ Other backends (and their limitations)
 Numba
 -----
 
-NumPy random generator can be used with Numba backend.
+NumPy random generators can be natively used with the Numba backend.
 
 >>> rng = pytensor.shared(np.random.default_rng(123), name="randomstate_rng")
 >>> x = pt.random.normal(rng=rng)
@@ -692,7 +676,6 @@ Inner graphs:
     └─ *4-<Scalar(float32, shape=())> [id K] <Scalar(float32, shape=())>
  ← normal_rv{"(),()->()"}.1 [id G] <Scalar(float64, shape=())>
     └─ ···
-<ipykernel.iostream.OutStream at 0x7fa5d3a475e0>
 
 >>> print(numba_fn(), numba_fn())
 -0.9891213503478509 -0.9891213503478509
@@ -700,11 +683,11 @@ Inner graphs:
 JAX
 ---
 
-JAX uses a different type of PRNG than those of Numpy. This means that the standard shared RNGs cannot be used directly in graphs transpiled to JAX.
+JAX uses a different type of PRNG than those of NumPy. This means that the standard shared RNGs cannot be used directly in graphs transpiled to JAX.
 
-Instead a copy of the Shared RNG variable is made, and its bit generator state is given a jax_state entry that is actually used by the JAX random variables.
+Instead a copy of the Shared RNG variable is made, and its bit generator state is expanded with a jax_state entry. This is what's actually used by the JAX random variables.
 
-In general, update rules are still respected, but they won't be used on the original shared variable, only the copied one actually used in the transpiled function
+In general, update rules are still respected, but they won't update/rely on the original shared variable.
 
 >>> import jax
 >>> rng = pytensor.shared(np.random.default_rng(123), name="rng")
@@ -718,7 +701,6 @@ uniform_rv{"(),()->()"}.1 [id A] <Scalar(float64, shape=())> 0
  └─ 1.0 [id E] <Scalar(float32, shape=())>
 uniform_rv{"(),()->()"}.0 [id A] <RandomGeneratorType> 0
  └─ ···
-<ipykernel.iostream.OutStream at 0x7fa5d3a475e0>
 
 >>> print(jax_fn(), jax_fn())
 [Array(0.07577298, dtype=float64)] [Array(0.09217023, dtype=float64)]

{pytensor-2.27.0 → pytensor-2.27.1}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-01-24T10:30:26+0100",
+ "date": "2025-02-03T19:06:21+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "4ea425938a3f82604d234ac1db2c3239b8d9cd8d",
- "version": "2.27.0"
+ "full-revisionid": "c22e79e1bbf54deb7043a3a778b6e91295d25e07",
+ "version": "2.27.1"
 }
 '''  # END VERSION_JSON
 

pytensor-2.27.1/pytensor/graph/rewriting/__init__.py

@@ -0,0 +1,4 @@
+from pytensor.graph.rewriting.utils import rewrite_graph
+
+
+all = ("rewrite_graph",)

{pytensor-2.27.0 → pytensor-2.27.1}/pytensor/graph/rewriting/db.py

@@ -35,6 +35,7 @@ class RewriteDatabase:
         rewriter: Union["RewriteDatabase", RewritesType],
         *tags: str,
         use_db_name_as_tag=True,
+        overwrite_existing=False,
     ):
         """Register a new rewriter to the database.
 
@@ -56,7 +57,8 @@ class RewriteDatabase:
             ``local_remove_all_assert``. Setting `use_db_name_as_tag` to
             ``False`` removes that behavior. This means that only the rewrite's name
             and/or its tags will enable it.
-
+        overwrite_existing:
+            Overwrite the existing rewriter with a new one having the same name
         """
         if not isinstance(
             rewriter,
@@ -66,22 +68,27 @@ class RewriteDatabase:
         ):
             raise TypeError(f"{rewriter} is not a valid rewrite type.")
 
-        if name in self.__db__:
-            raise ValueError(f"The tag '{name}' is already present in the database.")
-
         if use_db_name_as_tag:
             if self.name is not None:
                 tags = (*tags, self.name)
 
         rewriter.name = name
-        # This restriction is there because in many place we suppose that
-        # something in the RewriteDatabase is there only once.
-        if rewriter.name in self.__db__:
-            raise ValueError(
-                f"Tried to register {rewriter.name} again under the new name {name}. "
-                "The same rewrite cannot be registered multiple times in"
-                " an `RewriteDatabase`; use `ProxyDB` instead."
-            )
+
+        # if tag collides with name
+        if name in self.__db__ and name not in self._names:
+            raise ValueError(f"The tag '{name}' is already present in the database.")
+
+        if name in self.__db__ or rewriter.name in self.__db__:
+            if overwrite_existing:
+                self.remove_tags(name, *tags)
+                old_rewriter = self.__db__[name].pop()
+                self._names.remove(name)
+                self.__db__[old_rewriter.__class__.__name__].remove(old_rewriter)
+            else:
+                raise ValueError(
+                    f"The tag '{name}' is already present in the database."
+                )
+
         self.__db__[name] = OrderedSet([rewriter])
         self._names.add(name)
         self.__db__[rewriter.__class__.__name__].add(rewriter)