Repository: pymc-devs/pytensor

Issue:
Scan requires random variables to be referenced as non sequences
Works:

```python
import pytensor
import pytensor.tensor as at
import numpy as np

data = at.constant(np.random.randn(64))
srng = pytensor.tensor.random.RandomStream()
index = srng.integers(64, size=10)
datai = data[index]
var = at.vector("var")
scan = pytensor.scan(
    lambda v, *_: ((datai-v)**2).sum(), 
    sequences=var, non_sequences=[index], 
    strict=True
)
print(scan[0].eval({var: np.array([1., 2.])}))
print(pytensor.grad(scan[0].sum(), var).eval({var: [1, 1]}))
```

Raises an uninformative or fairly informative error
```python
import pytensor
import pytensor.tensor as at
import numpy as np

data = at.constant(np.random.randn(64))
srng = pytensor.tensor.random.RandomStream()
index = srng.integers(64, size=10)
datai = data[index]
var = at.vector("var")
scan = pytensor.scan(
    lambda v: ((datai-v)**2).sum(), 
    sequences=var, 
    strict=True
)
print(scan[0].eval({var: np.array([1., 2.])}))
print(pytensor.grad(scan[0].sum(), var).eval({var: [1, 1]}))
```

<details>

```
---------------------------------------------------------------------------
MissingInputError                         Traceback (most recent call last)
Cell In [129], line 10
      8 datai = data[index]
      9 var = at.vector("var")
---> 10 scan = aesara.scan(
     11     lambda v, *_: ((datai-v)**2).sum(), 
     12     sequences=var, #non_sequences=[index],
     13     strict=True
     14 )
     15 print(scan[0].eval({var: np.array([1., 2.])}))
     16 print(aesara.grad(scan[0].sum(), var).eval({var: [1, 1]}))

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/scan/basic.py:1140, in scan(fn, sequences, outputs_info, non_sequences, n_steps, truncate_gradient, go_backwards, mode, name, profile, allow_gc, strict, return_list)
   1126     allow_gc = config.scan__allow_gc
   1128 info = ScanInfo(
   1129     n_seqs=n_seqs,
   1130     mit_mot_in_slices=(),
   (...)
   1137     as_while=as_while,
   1138 )
-> 1140 local_op = Scan(
   1141     inner_inputs,
   1142     new_outs,
   1143     info,
   1144     mode=mode,
   1145     truncate_gradient=truncate_gradient,
   1146     name=name,
   1147     profile=profile,
   1148     allow_gc=allow_gc,
   1149     strict=strict,
   1150 )
   1152 ##
   1153 # Step 8. Compute the outputs using the scan op
   1154 ##
   1155 _scan_inputs = (
   1156     scan_seqs
   1157     + mit_mot_scan_inputs
   (...)
   1163     + other_scan_args
   1164 )

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/scan/op.py:859, in Scan.__init__(self, inputs, outputs, info, mode, typeConstructor, truncate_gradient, name, as_while, profile, allow_gc, strict)
    856 self.n_outer_inputs = info.n_outer_inputs
    857 self.n_outer_outputs = info.n_outer_outputs
--> 859 self.fgraph = FunctionGraph(inputs, outputs, clone=False)
    861 _ = self.prepare_fgraph(self.fgraph)
    863 if any(node.op.destroy_map for node in self.fgraph.apply_nodes):

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/graph/fg.py:153, in FunctionGraph.__init__(self, inputs, outputs, features, clone, update_mapping, **clone_kwds)
    150     self.add_input(in_var, check=False)
    152 for output in outputs:
--> 153     self.add_output(output, reason="init")
    155 self.profile = None
    156 self.update_mapping = update_mapping

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/graph/fg.py:163, in FunctionGraph.add_output(self, var, reason, import_missing)
    161 """Add a new variable as an output to this `FunctionGraph`."""
    162 self.outputs.append(var)
--> 163 self.import_var(var, reason=reason, import_missing=import_missing)
    164 self.clients[var].append(("output", len(self.outputs) - 1))

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/graph/fg.py:304, in FunctionGraph.import_var(self, var, reason, import_missing)
    302 # Imports the owners of the variables
    303 if var.owner and var.owner not in self.apply_nodes:
--> 304     self.import_node(var.owner, reason=reason, import_missing=import_missing)
    305 elif (
    306     var.owner is None
    307     and not isinstance(var, AtomicVariable)
    308     and var not in self.inputs
    309 ):
    310     from aesara.graph.null_type import NullType

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/graph/fg.py:369, in FunctionGraph.import_node(self, apply_node, check, reason, import_missing)
    360                 else:
    361                     error_msg = (
    362                         f"Input {node.inputs.index(var)} ({var})"
    363                         " of the graph (indices start "
   (...)
    367                         "for more information on this error."
    368                     )
--> 369                     raise MissingInputError(error_msg, variable=var)
    371 for node in new_nodes:
    372     assert node not in self.apply_nodes

MissingInputError: Input 0 (RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FCBAFBA03C0>)) of the graph (indices start from 0), used to compute integers_rv{0, (0, 0), int64, False}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FCBAFBA03C0>), TensorConstant{(1,) of 10}, TensorConstant{4}, TensorConstant{0}, TensorConstant{64}), was not provided and not given a value. Use the Aesara flag exception_verbosity='high', for more information on this error.
 
Backtrace when that variable is created:

  File "/home/mkochurov/micromamba/envs/bayes/lib/python3.9/site-packages/ipykernel/zmqshell.py", line 528, in run_cell
    return super().run_cell(*args, **kwargs)
  File "/home/mkochurov/micromamba/envs/bayes/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 2940, in run_cell
    result = self._run_cell(
  File "/home/mkochurov/micromamba/envs/bayes/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 2995, in _run_cell
    return runner(coro)
  File "/home/mkochurov/micromamba/envs/bayes/lib/python3.9/site-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner
    coro.send(None)
  File "/home/mkochurov/micromamba/envs/bayes/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3194, in run_cell_async
    has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
  File "/home/mkochurov/micromamba/envs/bayes/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3373, in run_ast_nodes
    if await self.run_code(code, result, async_=asy):
  File "/home/mkochurov/micromamba/envs/bayes/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3433, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "/tmp/ipykernel_441370/4148489949.py", line 7, in <module>
    index = srng.integers(64, size=10)
```
</details>

Scan silently passes and aesara.grad fails miserably later if strict=False

```python
import pytensor
import pytensor.tensor as at
import numpy as np

data = at.constant(np.random.randn(64))
srng = pytensor.tensor.random.RandomStream()
index = srng.integers(64, size=10)
datai = data[index]
var = at.vector("var")
scan = pytensor.scan(
    lambda v: ((datai-v)**2).sum(), 
    sequences=var 
    strict=False
)
print(scan[0].eval({var: np.array([1., 2.])}))
print(pytensor.grad(scan[0].sum(), var).eval({var: [1, 1]}))
```

<details>

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In [130], line 16
     10 scan = aesara.scan(
     11     lambda v, *_: ((datai-v)**2).sum(), 
     12     sequences=var, #non_sequences=[index],
     13     strict=False
     14 )
     15 print(scan[0].eval({var: np.array([1., 2.])}))
---> 16 print(aesara.grad(scan[0].sum(), var).eval({var: [1, 1]}))

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/gradient.py:623, in grad(cost, wrt, consider_constant, disconnected_inputs, add_names, known_grads, return_disconnected, null_gradients)
    620     if hasattr(g.type, "dtype"):
    621         assert g.type.dtype in aesara.tensor.type.float_dtypes
--> 623 _rval: Sequence[Variable] = _populate_grad_dict(
    624     var_to_app_to_idx, grad_dict, _wrt, cost_name
    625 )
    627 rval: MutableSequence[Optional[Variable]] = list(_rval)
    629 for i in range(len(_rval)):

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/gradient.py:1434, in _populate_grad_dict(var_to_app_to_idx, grad_dict, wrt, cost_name)
   1431     # end if cache miss
   1432     return grad_dict[var]
-> 1434 rval = [access_grad_cache(elem) for elem in wrt]
   1436 return rval

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/gradient.py:1434, in <listcomp>(.0)
   1431     # end if cache miss
   1432     return grad_dict[var]
-> 1434 rval = [access_grad_cache(elem) for elem in wrt]
   1436 return rval

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/gradient.py:1387, in _populate_grad_dict.<locals>.access_grad_cache(var)
   1384 for node in node_to_idx:
   1385     for idx in node_to_idx[node]:
-> 1387         term = access_term_cache(node)[idx]
   1389         if not isinstance(term, Variable):
   1390             raise TypeError(
   1391                 f"{node.op}.grad returned {type(term)}, expected"
   1392                 " Variable instance."
   1393             )

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/gradient.py:1058, in _populate_grad_dict.<locals>.access_term_cache(node)
   1054 if node not in term_dict:
   1056     inputs = node.inputs
-> 1058     output_grads = [access_grad_cache(var) for var in node.outputs]
   1060     # list of bools indicating if each output is connected to the cost
   1061     outputs_connected = [
   1062         not isinstance(g.type, DisconnectedType) for g in output_grads
   1063     ]

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/gradient.py:1058, in <listcomp>(.0)
   1054 if node not in term_dict:
   1056     inputs = node.inputs
-> 1058     output_grads = [access_grad_cache(var) for var in node.outputs]
   1060     # list of bools indicating if each output is connected to the cost
   1061     outputs_connected = [
   1062         not isinstance(g.type, DisconnectedType) for g in output_grads
   1063     ]

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/gradient.py:1387, in _populate_grad_dict.<locals>.access_grad_cache(var)
   1384 for node in node_to_idx:
   1385     for idx in node_to_idx[node]:
-> 1387         term = access_term_cache(node)[idx]
   1389         if not isinstance(term, Variable):
   1390             raise TypeError(
   1391                 f"{node.op}.grad returned {type(term)}, expected"
   1392                 " Variable instance."
   1393             )

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/gradient.py:1058, in _populate_grad_dict.<locals>.access_term_cache(node)
   1054 if node not in term_dict:
   1056     inputs = node.inputs
-> 1058     output_grads = [access_grad_cache(var) for var in node.outputs]
   1060     # list of bools indicating if each output is connected to the cost
   1061     outputs_connected = [
   1062         not isinstance(g.type, DisconnectedType) for g in output_grads
   1063     ]

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/gradient.py:1058, in <listcomp>(.0)
   1054 if node not in term_dict:
   1056     inputs = node.inputs
-> 1058     output_grads = [access_grad_cache(var) for var in node.outputs]
   1060     # list of bools indicating if each output is connected to the cost
   1061     outputs_connected = [
   1062         not isinstance(g.type, DisconnectedType) for g in output_grads
   1063     ]

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/gradient.py:1387, in _populate_grad_dict.<locals>.access_grad_cache(var)
   1384 for node in node_to_idx:
   1385     for idx in node_to_idx[node]:
-> 1387         term = access_term_cache(node)[idx]
   1389         if not isinstance(term, Variable):
   1390             raise TypeError(
   1391                 f"{node.op}.grad returned {type(term)}, expected"
   1392                 " Variable instance."
   1393             )

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/gradient.py:1213, in _populate_grad_dict.<locals>.access_term_cache(node)
   1205         if o_shape != g_shape:
   1206             raise ValueError(
   1207                 "Got a gradient of shape "
   1208                 + str(o_shape)
   1209                 + " on an output of shape "
   1210                 + str(g_shape)
   1211             )
-> 1213 input_grads = node.op.L_op(inputs, node.outputs, new_output_grads)
   1215 if input_grads is None:
   1216     raise TypeError(
   1217         f"{node.op}.grad returned NoneType, expected iterable."
   1218     )

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/scan/op.py:2613, in Scan.L_op(self, inputs, outs, dC_douts)
   2611 for dx in range(len(dC_dinps_t)):
   2612     if not dC_dinps_t[dx]:
-> 2613         dC_dinps_t[dx] = at.zeros_like(diff_inputs[dx])
   2614     else:
   2615         disconnected_dC_dinps_t[dx] = False

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/tensor/basic.py:798, in zeros_like(model, dtype, opt)
    782 def zeros_like(model, dtype=None, opt=False):
    783     """equivalent of numpy.zeros_like
    784     Parameters
    785     ----------
   (...)
    795         tensor the shape of model containing zeros of the type of dtype.
    796     """
--> 798     _model = as_tensor_variable(model)
    800     if dtype is None:
    801         dtype = _model.type.dtype

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/tensor/__init__.py:49, in as_tensor_variable(x, name, ndim, **kwargs)
     17 def as_tensor_variable(
     18     x: TensorLike, name: Optional[str] = None, ndim: Optional[int] = None, **kwargs
     19 ) -> "TensorVariable":
     20     """Convert `x` into an equivalent `TensorVariable`.
     21 
     22     This function can be used to turn ndarrays, numbers, `ScalarType` instances,
   (...)
     47 
     48     """
---> 49     return _as_tensor_variable(x, name, ndim, **kwargs)

File ~/micromamba/envs/bayes/lib/python3.9/functools.py:888, in singledispatch.<locals>.wrapper(*args, **kw)
    884 if not args:
    885     raise TypeError(f'{funcname} requires at least '
    886                     '1 positional argument')
--> 888 return dispatch(args[0].__class__)(*args, **kw)

File ~/micromamba/envs/bayes/lib/python3.9/site-packages/aesara/tensor/basic.py:100, in _as_tensor_Variable(x, name, ndim, **kwargs)
     97 @_as_tensor_variable.register(Variable)
     98 def _as_tensor_Variable(x, name, ndim, **kwargs):
     99     if not isinstance(x.type, TensorType):
--> 100         raise TypeError(
    101             f"Tensor type field must be a TensorType; found {type(x.type)}."
    102         )
    104     if ndim is None:
    105         return x

TypeError: Tensor type field must be a TensorType; found <class 'aesara.tensor.random.type.RandomGeneratorType'>.
```

</details>
**Please provide any additional information below.**


## Versions and main components

* Aesara version: '2.8.7'
* Python version: '3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:58:50) [GCC 10.3.0]'
* Operating system: Ubuntu
* How did you install Aesara: (conda/pip) conda

<details> <summary> pytensor config: </summary>

```
floatX ({'float64', 'float16', 'float32'}) 
    Doc:  Default floating-point precision for python casts.

Note: float16 support is experimental, use at your own risk.
    Value:  float64

warn_float64 ({'ignore', 'pdb', 'warn', 'raise'}) 
    Doc:  Do an action when a tensor variable with float64 dtype is created.
    Value:  ignore

pickle_test_value (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcc2c19ac10>>) 
    Doc:  Dump test values while pickling model. If True, test values will be dumped with model.
    Value:  True

cast_policy ({'custom', 'numpy+floatX'}) 
    Doc:  Rules for implicit type casting
    Value:  custom

deterministic ({'more', 'default'}) 
    Doc:  If `more`, sometimes we will select some implementation that are more deterministic, but slower.  Also see the dnn.conv.algo* flags to cover more cases.
    Value:  default

device (cpu)
    Doc:  Default device for computations. only cpu is supported for now
    Value:  cpu

force_device (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcc2c19a610>>) 
    Doc:  Raise an error if we can't use the specified device
    Value:  False

conv__assert_shape (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcc2c19a5b0>>) 
    Doc:  If True, AbstractConv* ops will verify that user-provided shapes match the runtime shapes (debugging option, may slow down compilation)
    Value:  False

print_global_stats (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5baa670>>) 
    Doc:  Print some global statistics (time spent) at the end
    Value:  False

assert_no_cpu_op ({'ignore', 'pdb', 'warn', 'raise'}) 
    Doc:  Raise an error/warning if there is a CPU op in the computational graph.
    Value:  ignore

unpickle_function (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5baa850>>) 
    Doc:  Replace unpickled Aesara functions with None. This is useful to unpickle old graphs that pickled them when it shouldn't
    Value:  True

<aesara.configparser.ConfigParam object at 0x7fcbc5baa8b0>
    Doc:  Default compilation mode
    Value:  Mode

cxx (<class 'str'>) 
    Doc:  The C++ compiler to use. Currently only g++ is supported, but supporting additional compilers should not be too difficult. If it is empty, no C++ code is compiled.
    Value:  /usr/bin/g++

linker ({'vm_nogc', 'c|py_nogc', 'c', 'py', 'c|py', 'vm', 'cvm_nogc', 'cvm'}) 
    Doc:  Default linker used if the aesara flags mode is Mode
    Value:  cvm

allow_gc (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5baaa30>>) 
    Doc:  Do we default to delete intermediate results during Aesara function calls? Doing so lowers the memory requirement, but asks that we reallocate memory at the next function call. This is implemented for the default linker, but may not work for all linkers.
    Value:  True

optimizer ({'o2', 'o1', 'o3', 'None', 'fast_compile', 'o4', 'fast_run', 'unsafe', 'merge'}) 
    Doc:  Default optimizer. If not None, will use this optimizer with the Mode
    Value:  o4

optimizer_verbose (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41550>>) 
    Doc:  If True, we print all optimization being applied
    Value:  False

on_opt_error ({'ignore', 'pdb', 'warn', 'raise'}) 
    Doc:  What to do when an optimization crashes: warn and skip it, raise the exception, or fall into the pdb debugger.
    Value:  warn

nocleanup (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b415b0>>) 
    Doc:  Suppress the deletion of code files that did not compile cleanly
    Value:  False

on_unused_input ({'ignore', 'warn', 'raise'}) 
    Doc:  What to do if a variable in the 'inputs' list of  aesara.function() is not used in the graph.
    Value:  raise

gcc__cxxflags (<class 'str'>) 
    Doc:  Extra compiler flags for gcc
    Value:   -Wno-c++11-narrowing -fno-exceptions -fno-unwind-tables -fno-asynchronous-unwind-tables

cmodule__warn_no_version (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41640>>) 
    Doc:  If True, will print a warning when compiling one or more Op with C code that can't be cached because there is no c_code_cache_version() function associated to at least one of those Ops.
    Value:  False

cmodule__remove_gxx_opt (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41400>>) 
    Doc:  If True, will remove the -O* parameter passed to g++.This is useful to debug in gdb modules compiled by Aesara.The parameter -g is passed by default to g++
    Value:  False

cmodule__compilation_warning (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41670>>) 
    Doc:  If True, will print compilation warnings.
    Value:  False

cmodule__preload_cache (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b416a0>>) 
    Doc:  If set to True, will preload the C module cache at import time
    Value:  False

cmodule__age_thresh_use (<class 'int'>) 
    Doc:  In seconds. The time after which Aesara won't reuse a compile c module.
    Value:  2073600

cmodule__debug (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41700>>) 
    Doc:  If True, define a DEBUG macro (if not exists) for any compiled C code.
    Value:  False

compile__wait (<class 'int'>) 
    Doc:  Time to wait before retrying to acquire the compile lock.
    Value:  5

compile__timeout (<class 'int'>) 
    Doc:  In seconds, time that a process will wait before deciding to
    override an existing lock. An override only happens when the existing
    lock is held by the same owner *and* has not been 'refreshed' by this
    owner for more than this period. Refreshes are done every half timeout
    period for running processes.
    Value:  120

ctc__root (<class 'str'>) 
    Doc:  Directory which contains the root of Baidu CTC library. It is assumed         that the compiled library is either inside the build, lib or lib64         subdirectory, and the header inside the include directory.
    Value:  

tensor__cmp_sloppy (<class 'int'>) 
    Doc:  Relax aesara.tensor.math._allclose (0) not at all, (1) a bit, (2) more
    Value:  0

tensor__local_elemwise_fusion (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b418b0>>) 
    Doc:  Enable or not in fast_run mode(fast_run optimization) the elemwise fusion optimization
    Value:  True

lib__amblibm (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41940>>) 
    Doc:  Use amd's amdlibm numerical library
    Value:  False

tensor__insert_inplace_optimizer_validate_nb (<class 'int'>) 
    Doc:  -1: auto, if graph have less then 500 nodes 1, else 10
    Value:  -1

traceback__limit (<class 'int'>) 
    Doc:  The number of stack to trace. -1 mean all.
    Value:  8

traceback__compile_limit (<class 'int'>) 
    Doc:  The number of stack to trace to keep during compilation. -1 mean all. If greater then 0, will also make us save Aesara internal stack trace.
    Value:  0

experimental__local_alloc_elemwise (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41ac0>>) 
    Doc:  DEPRECATED: If True, enable the experimental optimization local_alloc_elemwise. Generates error if not True. Use optimizer_excluding=local_alloc_elemwise to disable.
    Value:  True

experimental__local_alloc_elemwise_assert (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41af0>>) 
    Doc:  When the local_alloc_elemwise is applied, add an assert to highlight shape errors.
    Value:  True

warn__ignore_bug_before ({'0.8.2', '0.10', 'all', '0.5', '0.7', '1.0.3', '1.0.5', 'None', '0.8', '1.0.2', '0.3', '1.0', '0.4', '0.6', '0.9', '0.4.1', '1.0.1', '0.8.1', '1.0.4'}) 
    Doc:  If 'None', we warn about all Aesara bugs found by default. If 'all', we don't warn about Aesara bugs found by default. If a version, we print only the warnings relative to Aesara bugs found after that version. Warning for specific bugs can be configured with specific [warn] flags.
    Value:  0.9

exception_verbosity ({'high', 'low'}) 
    Doc:  If 'low', the text of exceptions will generally refer to apply nodes with short names such as Elemwise{add_no_inplace}. If 'high', some exceptions will also refer to apply nodes with long descriptions  like:
        A. Elemwise{add_no_inplace}
                B. log_likelihood_v_given_h
                C. log_likelihood_h
    Value:  low

print_test_value (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41c40>>) 
    Doc:  If 'True', the __eval__ of an Aesara variable will return its test_value when this is available. This has the practical conseguence that, e.g., in debugging `my_var` will print the same as `my_var.tag.test_value` when a test value is defined.
    Value:  False

compute_test_value ({'raise', 'ignore', 'pdb', 'warn', 'off'}) 
    Doc:  If 'True', Aesara will run each op at graph build time, using Constants, SharedVariables and the tag 'test_value' as inputs to the function. This helps the user track down problems in the graph before it gets optimized.
    Value:  off

compute_test_value_opt ({'raise', 'ignore', 'pdb', 'warn', 'off'}) 
    Doc:  For debugging Aesara optimization only. Same as compute_test_value, but is used during Aesara optimization
    Value:  off

check_input (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41cd0>>) 
    Doc:  Specify if types should check their input in their C code. It can be used to speed up compilation, reduce overhead (particularly for scalars) and reduce the number of generated C files.
    Value:  True

NanGuardMode__nan_is_error (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41d00>>) 
    Doc:  Default value for nan_is_error
    Value:  True

NanGuardMode__inf_is_error (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41d30>>) 
    Doc:  Default value for inf_is_error
    Value:  True

NanGuardMode__big_is_error (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41dc0>>) 
    Doc:  Default value for big_is_error
    Value:  True

NanGuardMode__action ({'pdb', 'warn', 'raise'}) 
    Doc:  What NanGuardMode does when it finds a problem
    Value:  raise

DebugMode__patience (<class 'int'>) 
    Doc:  Optimize graph this many times to detect inconsistency
    Value:  10

DebugMode__check_c (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41e50>>) 
    Doc:  Run C implementations where possible
    Value:  True

DebugMode__check_py (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41ee0>>) 
    Doc:  Run Python implementations where possible
    Value:  True

DebugMode__check_finite (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41f10>>) 
    Doc:  True -> complain about NaN/Inf results
    Value:  True

DebugMode__check_strides (<class 'int'>) 
    Doc:  Check that Python- and C-produced ndarrays have same strides. On difference: (0) - ignore, (1) warn, or (2) raise error
    Value:  0

DebugMode__warn_input_not_reused (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41f70>>) 
    Doc:  Generate a warning when destroy_map or view_map says that an op works inplace, but the op did not reuse the input for its output.
    Value:  True

DebugMode__check_preallocated_output (<class 'str'>) 
    Doc:  Test thunks with pre-allocated memory as output storage. This is a list of strings separated by ":". Valid values are: "initial" (initial storage in storage map, happens with Scan),"previous" (previously-returned memory), "c_contiguous", "f_contiguous", "strided" (positive and negative strides), "wrong_size" (larger and smaller dimensions), and "ALL" (all of the above).
    Value:  

DebugMode__check_preallocated_output_ndim (<class 'int'>) 
    Doc:  When testing with "strided" preallocated output memory, test all combinations of strides over that number of (inner-most) dimensions. You may want to reduce that number to reduce memory or time usage, but it is advised to keep a minimum of 2.
    Value:  4

profiling__time_thunks (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b41e80>>) 
    Doc:  Time individual thunks when profiling
    Value:  True

profiling__n_apply (<class 'int'>) 
    Doc:  Number of Apply instances to print by default
    Value:  20

profiling__n_ops (<class 'int'>) 
    Doc:  Number of Ops to print by default
    Value:  20

profiling__output_line_width (<class 'int'>) 
    Doc:  Max line width for the profiling output
    Value:  512

profiling__min_memory_size (<class 'int'>) 
    Doc:  For the memory profile, do not print Apply nodes if the size
                 of their outputs (in bytes) is lower than this threshold
    Value:  1024

profiling__min_peak_memory (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b52190>>) 
    Doc:  The min peak memory usage of the order
    Value:  False

profiling__destination (<class 'str'>) 
    Doc:  File destination of the profiling output
    Value:  stderr

profiling__debugprint (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b521f0>>) 
    Doc:  Do a debugprint of the profiled functions
    Value:  False

profiling__ignore_first_call (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b52220>>) 
    Doc:  Do we ignore the first call of an Aesara function.
    Value:  False

on_shape_error ({'warn', 'raise'}) 
    Doc:  warn: print a warning and use the default value. raise: raise an error
    Value:  warn

openmp (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b52280>>) 
    Doc:  Allow (or not) parallel computation on the CPU with OpenMP. This is the default value used when creating an Op that supports OpenMP parallelization. It is preferable to define it via the Aesara configuration file ~/.aesararc or with the environment variable AESARA_FLAGS. Parallelization is only done for some operations that implement it, and even for operations that implement parallelism, each operation is free to respect this flag or not. You can control the number of threads used with the environment variable OMP_NUM_THREADS. If it is set to 1, we disable openmp in Aesara by default.
    Value:  False

openmp_elemwise_minsize (<class 'int'>) 
    Doc:  If OpenMP is enabled, this is the minimum size of vectors for which the openmp parallelization is enabled in element wise ops.
    Value:  200000

optimizer_excluding (<class 'str'>) 
    Doc:  When using the default mode, we will remove optimizer with these tags. Separate tags with ':'.
    Value:  

optimizer_including (<class 'str'>) 
    Doc:  When using the default mode, we will add optimizer with these tags. Separate tags with ':'.
    Value:  

optimizer_requiring (<class 'str'>) 
    Doc:  When using the default mode, we will require optimizer with these tags. Separate tags with ':'.
    Value:  

optdb__position_cutoff (<class 'float'>) 
    Doc:  Where to stop eariler during optimization. It represent the position of the optimizer where to stop.
    Value:  inf

optdb__max_use_ratio (<class 'float'>) 
    Doc:  A ratio that prevent infinite loop in EquilibriumGraphRewriter.
    Value:  8.0

cycle_detection ({'fast', 'regular'}) 
    Doc:  If cycle_detection is set to regular, most inplaces are allowed,but it is slower. If cycle_detection is set to faster, less inplacesare allowed, but it makes the compilation faster.The interaction of which one give the lower peak memory usage iscomplicated and not predictable, so if you are close to the peakmemory usage, triyng both could give you a small gain.
    Value:  regular

check_stack_trace ({'log', 'off', 'warn', 'raise'}) 
    Doc:  A flag for checking the stack trace during the optimization process. default (off): does not check the stack trace of any optimization log: inserts a dummy stack trace that identifies the optimizationthat inserted the variable that had an empty stack trace.warn: prints a warning if a stack trace is missing and also a dummystack trace is inserted that indicates which optimization insertedthe variable that had an empty stack trace.raise: raises an exception if a stack trace is missing
    Value:  off

metaopt__verbose (<class 'int'>) 
    Doc:  0 for silent, 1 for only warnings, 2 for full output withtimings and selected implementation
    Value:  0

metaopt__optimizer_excluding (<class 'str'>) 
    Doc:  exclude optimizers with these tags. Separate tags with ':'.
    Value:  

metaopt__optimizer_including (<class 'str'>) 
    Doc:  include optimizers with these tags. Separate tags with ':'.
    Value:  

profile (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b52580>>) 
    Doc:  If VM should collect profile information
    Value:  False

profile_optimizer (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b525b0>>) 
    Doc:  If VM should collect optimizer profile information
    Value:  False

profile_memory (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b525e0>>) 
    Doc:  If VM should collect memory profile information and print it
    Value:  False

<aesara.configparser.ConfigParam object at 0x7fcbc5b52610>
    Doc:  Useful only for the VM Linkers. When lazy is None, auto detect if lazy evaluation is needed and use the appropriate version. If the C loop isn't being used and lazy is True, use the Stack VM; otherwise, use the Loop VM.
    Value:  None

unittests__rseed (<class 'str'>) 
    Doc:  Seed to use for randomized unit tests. Special value 'random' means using a seed of None.
    Value:  666

warn__round (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b526d0>>) 
    Doc:  Warn when using `tensor.round` with the default mode. Round changed its default from `half_away_from_zero` to `half_to_even` to have the same default as NumPy.
    Value:  False

numba__vectorize_target ({'cuda', 'parallel', 'cpu'}) 
    Doc:  Default target for numba.vectorize.
    Value:  cpu

numba__fastmath (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b52790>>) 
    Doc:  If True, use Numba's fastmath mode.
    Value:  True

numba__cache (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbc5b52820>>) 
    Doc:  If True, use Numba's file based caching.
    Value:  True

compiledir_format (<class 'str'>) 
    Doc:  Format string for platform-dependent compiled module subdirectory
(relative to base_compiledir). Available keys: aesara_version, device,
gxx_version, hostname, numpy_version, platform, processor,
python_bitwidth, python_int_bitwidth, python_version, short_platform.
Defaults to compiledir_%(short_platform)s-%(processor)s-%(python_versi
on)s-%(python_bitwidth)s.
    Value:  compiledir_%(short_platform)s-%(processor)s-%(python_version)s-%(python_bitwidth)s

<aesara.configparser.ConfigParam object at 0x7fcbc5b528e0>
    Doc:  platform-independent root directory for compiled modules
    Value:  /home/mkochurov/.aesara

<aesara.configparser.ConfigParam object at 0x7fcbc5b527f0>
    Doc:  platform-dependent cache directory for compiled modules
    Value:  /home/mkochurov/.aesara/compiledir_Linux-5.4--generic-x86_64-with-glibc2.31-x86_64-3.9.13-64

blas__ldflags (<class 'str'>) 
    Doc:  lib[s] to include for [Fortran] level-3 blas implementation
    Value:  -L/home/mkochurov/micromamba/envs/bayes/lib -lmkl_core -lmkl_intel_thread -lmkl_rt -Wl,-rpath,/home/mkochurov/micromamba/envs/bayes/lib

blas__check_openmp (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcc05b91640>>) 
    Doc:  Check for openmp library conflict.
WARNING: Setting this to False leaves you open to wrong results in blas-related operations.
    Value:  True

scan__allow_gc (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcbba30df10>>) 
    Doc:  Allow/disallow gc inside of Scan (default: False)
    Value:  False

scan__allow_output_prealloc (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fcc3034d1c0>>) 
    Doc:  Allow/disallow memory preallocation for outputs inside of scan (default: True)
    Value:  True
```

</details>

Hints from issue comments:
@ricardoV94: Scan and OpFromGraph are a bit too eager in how they handle SharedVariables, even if they wouldn't necessarily be needed as inputs of the inner graph. For Scan at least, this is a bit related to the whole updates machinery, so that RandomVariables update across iterations without users saying so manually...

Solution patch:
diff --git a/doc/extending/op.rst b/doc/extending/op.rst
index ddd397dee..b1585c4ec 100644
--- a/doc/extending/op.rst
+++ b/doc/extending/op.rst
@@ -506,4 +506,3 @@ These are the function required to work with :func:`pytensor.gradient.grad`.
    the outputs) back to their corresponding shapes and return them as the
    output of the :meth:`Op.R_op` method.
 
-   :ref:`List of op with r op support <R_op_list>`.
diff --git a/doc/library/gradient.rst b/doc/library/gradient.rst
deleted file mode 100644
index f823a1c38..000000000
--- a/doc/library/gradient.rst
+++ /dev/null
@@ -1,76 +0,0 @@
-.. _libdoc_gradient:
-
-===========================================
-:mod:`gradient` -- Symbolic Differentiation
-===========================================
-
-.. module:: gradient
-   :platform: Unix, Windows
-   :synopsis: low-level automatic differentiation
-.. moduleauthor:: LISA
-
-.. testsetup:: *
-
-   from pytensor.gradient import *
-
-Symbolic gradient is usually computed from :func:`gradient.grad`, which offers a
-more convenient syntax for the common case of wanting the gradient of some
-scalar cost with respect to some input expressions. The :func:`grad_sources_inputs`
-function does the underlying work, and is more flexible, but is also more
-awkward to use when :func:`gradient.grad` can do the job.
-
-
-Gradient related functions
-==========================
-
-.. automodule:: pytensor.gradient
-    :members:
-
-.. _R_op_list:
-
-
-List of Implemented R op
-========================
-
-
-See the :ref:`gradient tutorial <tutcomputinggrads>` for the R op documentation.
-
-list of ops that support R-op:
- * with test
-    * SpecifyShape
-    * MaxAndArgmax
-    * Subtensor
-    * IncSubtensor set_subtensor too
-    * Alloc
-    * Dot
-    * Elemwise
-    * Sum
-    * Softmax
-    * Shape
-    * Join
-    * Rebroadcast
-    * Reshape
-    * DimShuffle
-    * Scan [In tests/scan/test_basic.test_rop]
-
- * without test
-    * Split
-    * ARange
-    * ScalarFromTensor
-    * AdvancedSubtensor1
-    * AdvancedIncSubtensor1
-    * AdvancedIncSubtensor
-
-Partial list of ops without support for R-op:
-
- * All sparse ops
- * All linear algebra ops.
- * PermuteRowElements
- * AdvancedSubtensor
- * TensorDot
- * Outer
- * Prod
- * MulwithoutZeros
- * ProdWithoutZeros
- * CAReduce(for max,... done for MaxAndArgmax op)
- * MaxAndArgmax(only for matrix on axis 0 or 1)
diff --git a/doc/library/tensor/basic.rst b/doc/library/tensor/basic.rst
index 8d22c1e57..4f087b678 100644
--- a/doc/library/tensor/basic.rst
+++ b/doc/library/tensor/basic.rst
@@ -1791,5 +1791,3 @@ Gradient / Differentiation
     :members: grad
     :noindex:
 
-See the :ref:`gradient <libdoc_gradient>` page for complete documentation
-of the gradient module.
diff --git a/doc/tutorial/gradients.rst b/doc/tutorial/gradients.rst
index edb38bb01..f8b7f7ff9 100644
--- a/doc/tutorial/gradients.rst
+++ b/doc/tutorial/gradients.rst
@@ -86,9 +86,7 @@ of symbolic differentiation).
    ``i`` of the output list is the gradient of the first argument of
    `pt.grad` with respect to the ``i``-th element of the list given as second argument.
    The first argument of `pt.grad` has to be a scalar (a tensor
-   of size 1). For more information on the semantics of the arguments of
-   `pt.grad` and details about the implementation, see
-   :ref:`this<libdoc_gradient>` section of the library.
+   of size 1).
 
    Additional information on the inner workings of differentiation may also be
    found in the more advanced tutorial :ref:`Extending PyTensor<extending>`.
@@ -204,7 +202,21 @@ you need to do something similar to this:
 >>> f([[1, 1], [1, 1]], [[2, 2], [2, 2]], [0,1])
 array([ 2.,  2.])
 
-:ref:`List <R_op_list>` of Op that implement Rop.
+By default, the R-operator is implemented as a double application of the L_operator
+(see `reference <https://j-towns.github.io/2017/06/12/A-new-trick.html>`_).
+In most cases this should be as performant as a specialized implementation of the R-operator.
+However, PyTensor may sometimes fail to prune dead branches or fuse common expressions within composite operators,
+such as Scan and OpFromGraph, that would be more easily avoidable in a direct implentation of the R-operator.
+
+When this is a concern, it is possible to force `Rop` to use the specialized `Op.R_op` methods by passing
+`use_op_rop_implementation=True`. Note that this will fail if the graph contains `Op`s that don't implement this method.
+
+
+>>> JV = pytensor.gradient.Rop(y, W, V, use_op_rop_implementation=True)
+>>> f = pytensor.function([W, V, x], JV)
+>>> f([[1, 1], [1, 1]], [[2, 2], [2, 2]], [0,1])
+array([ 2.,  2.])
+
 
 L-operator
 ----------
@@ -234,7 +246,6 @@ array([[ 0.,  0.],
     as the input parameter, while the result of the R-operator has a shape similar
     to that of the output.
 
-   :ref:`List of op with r op support <R_op_list>`.
 
 Hessian times a Vector
 ======================
diff --git a/pytensor/compile/builders.py b/pytensor/compile/builders.py
index 49baa3bb2..a4a3d1840 100644
--- a/pytensor/compile/builders.py
+++ b/pytensor/compile/builders.py
@@ -340,6 +340,12 @@ class OpFromGraph(Op, HasInnerGraph):
             ``None``, this will be used as the connection_pattern for this
             :class:`Op`.
 
+        .. warning::
+
+            rop overrides is ignored when `pytensor.gradient.Rop` is called with
+            `use_op_rop_implementation=False` (default). In this case the Lop
+            is used twice to obtain a mathematically equivalent Rop.
+
         strict: bool, default False
             If true, it raises when any variables needed to compute the inner graph
             are not provided as explici inputs. This can only happen for graphs with
@@ -641,7 +647,12 @@ class OpFromGraph(Op, HasInnerGraph):
             return rop_overrides
 
         eval_points = [inp_t() for inp_t in self.input_types]
-        fn_rop = partial(Rop, wrt=inner_inputs, eval_points=eval_points)
+        fn_rop = partial(
+            Rop,
+            wrt=inner_inputs,
+            eval_points=eval_points,
+            use_op_rop_implementation=True,
+        )
 
         callable_args = (inner_inputs, eval_points)
         if rop_overrides is None:
diff --git a/pytensor/gradient.py b/pytensor/gradient.py
index 13ca94338..04572b29d 100644
--- a/pytensor/gradient.py
+++ b/pytensor/gradient.py
@@ -142,13 +142,50 @@ class DisconnectedType(Type):
 disconnected_type = DisconnectedType()
 
 
-def Rop(
-    f: Variable | Sequence[Variable],
-    wrt: Variable | Sequence[Variable],
-    eval_points: Variable | Sequence[Variable],
+def pushforward_through_pullback(
+    outputs: Sequence[Variable],
+    inputs: Sequence[Variable],
+    tangents: Sequence[Variable],
     disconnected_outputs: Literal["ignore", "warn", "raise"] = "raise",
     return_disconnected: Literal["none", "zero", "disconnected"] = "zero",
-) -> Variable | None | Sequence[Variable | None]:
+) -> Sequence[Variable | None]:
+    """Compute the pushforward (Rop) through two applications of a pullback (Lop) operation.
+
+    References
+    ----------
+    .. [1] J. Towns, "A new trick for calculating Jacobian vector products", 2017.
+           Available: https://j-towns.github.io/2017/06/12/A-new-trick.html
+
+    """
+    # Cotangents are just auxiliary variables that should be pruned from the final graph,
+    # but that would require a graph rewrite before the user tries to compile a pytensor function.
+    # To avoid trouble we use .zeros_like() instead of .type(), which does not create a new root variable.
+    cotangents = [out.zeros_like(dtype=config.floatX) for out in outputs]  # type: ignore
+
+    input_cotangents = Lop(
+        f=outputs,
+        wrt=inputs,
+        eval_points=cotangents,
+        disconnected_inputs=disconnected_outputs,
+        return_disconnected="zero",
+    )
+
+    return Lop(
+        f=input_cotangents,  # type: ignore
+        wrt=cotangents,
+        eval_points=tangents,
+        disconnected_inputs="ignore",
+        return_disconnected=return_disconnected,
+    )
+
+
+def _rop_legacy(
+    f: Sequence[Variable],
+    wrt: Sequence[Variable],
+    eval_points: Sequence[Variable],
+    disconnected_outputs: Literal["ignore", "warn", "raise"] = "raise",
+    return_disconnected: Literal["none", "zero", "disconnected"] = "zero",
+) -> Sequence[Variable | None]:
     """Computes the R-operator applied to `f` with respect to `wrt` at `eval_points`.
 
     Mathematically this stands for the Jacobian of `f` right multiplied by the
@@ -190,38 +227,6 @@ def Rop(
         If `f` is a list/tuple, then return a list/tuple with the results.
     """
 
-    if not isinstance(wrt, list | tuple):
-        _wrt: list[Variable] = [pytensor.tensor.as_tensor_variable(wrt)]
-    else:
-        _wrt = [pytensor.tensor.as_tensor_variable(x) for x in wrt]
-
-    if not isinstance(eval_points, list | tuple):
-        _eval_points: list[Variable] = [pytensor.tensor.as_tensor_variable(eval_points)]
-    else:
-        _eval_points = [pytensor.tensor.as_tensor_variable(x) for x in eval_points]
-
-    if not isinstance(f, list | tuple):
-        _f: list[Variable] = [pytensor.tensor.as_tensor_variable(f)]
-    else:
-        _f = [pytensor.tensor.as_tensor_variable(x) for x in f]
-
-    if len(_wrt) != len(_eval_points):
-        raise ValueError("`wrt` must be the same length as `eval_points`.")
-
-    # Check that each element of wrt corresponds to an element
-    # of eval_points with the same dimensionality.
-    for i, (wrt_elem, eval_point) in enumerate(zip(_wrt, _eval_points, strict=True)):
-        try:
-            if wrt_elem.type.ndim != eval_point.type.ndim:
-                raise ValueError(
-                    f"Elements {i} of `wrt` and `eval_point` have mismatched dimensionalities: "
-                    f"{wrt_elem.type.ndim} and {eval_point.type.ndim}"
-                )
-        except AttributeError:
-            # wrt_elem and eval_point don't always have ndim like random type
-            # Tensor, Sparse have the ndim attribute
-            pass
-
     seen_nodes: dict[Apply, Sequence[Variable]] = {}
 
     def _traverse(node):
@@ -237,8 +242,8 @@ def Rop(
         # inputs of the node
         local_eval_points = []
         for inp in inputs:
-            if inp in _wrt:
-                local_eval_points.append(_eval_points[_wrt.index(inp)])
+            if inp in wrt:
+                local_eval_points.append(eval_points[wrt.index(inp)])
             elif inp.owner is None:
                 try:
                     local_eval_points.append(inp.zeros_like())
@@ -292,13 +297,13 @@ def Rop(
     # end _traverse
 
     # Populate the dictionary
-    for out in _f:
+    for out in f:
         _traverse(out.owner)
 
     rval: list[Variable | None] = []
-    for out in _f:
-        if out in _wrt:
-            rval.append(_eval_points[_wrt.index(out)])
+    for out in f:
+        if out in wrt:
+            rval.append(eval_points[wrt.index(out)])
         elif (
             seen_nodes.get(out.owner, None) is None
             or seen_nodes[out.owner][out.owner.outputs.index(out)] is None
@@ -337,6 +342,116 @@ def Rop(
         else:
             rval.append(seen_nodes[out.owner][out.owner.outputs.index(out)])
 
+    return rval
+
+
+def Rop(
+    f: Variable | Sequence[Variable],
+    wrt: Variable | Sequence[Variable],
+    eval_points: Variable | Sequence[Variable],
+    disconnected_outputs: Literal["ignore", "warn", "raise"] = "raise",
+    return_disconnected: Literal["none", "zero", "disconnected"] = "zero",
+    use_op_rop_implementation: bool = False,
+) -> Variable | None | Sequence[Variable | None]:
+    """Computes the R-operator applied to `f` with respect to `wrt` at `eval_points`.
+
+    Mathematically this stands for the Jacobian of `f` right multiplied by the
+    `eval_points`.
+
+    By default, the R-operator is implemented as a double application of the L_operator [1]_.
+    In most cases this should be as performant as a specialized implementation of the R-operator.
+    However, PyTensor may sometimes fail to prune dead branches or fuse common expressions within composite operators,
+    such as Scan and OpFromGraph, that would be more easily avoidable in a direct implentation of the R-operator.
+
+    When this is a concern, it is possible to force `Rop` to use the specialized `Op.R_op` methods by passing
+    `use_op_rop_implementation=True`. Note that this will fail if the graph contains `Op`s that don't implement this method.
+
+    Parameters
+    ----------
+    f
+        The outputs of the computational graph to which the R-operator is
+        applied.
+    wrt
+        Variables for which the R-operator of `f` is computed.
+    eval_points
+        Points at which to evaluate each of the variables in `wrt`.
+    disconnected_outputs
+        Defines the behaviour if some of the variables in `f`
+        have no dependency on any of the variable in `wrt` (or if
+        all links are non-differentiable). The possible values are:
+
+        - ``'ignore'``: considers that the gradient on these parameters is zero.
+        - ``'warn'``: consider the gradient zero, and print a warning.
+        - ``'raise'``: raise `DisconnectedInputError`.
+
+    return_disconnected
+        - ``'zero'`` : If ``wrt[i]`` is disconnected, return value ``i`` will be
+          ``wrt[i].zeros_like()``.
+        - ``'none'`` : If ``wrt[i]`` is disconnected, return value ``i`` will be
+          ``None``
+        - ``'disconnected'`` : returns variables of type `DisconnectedType`
+    use_op_lop_implementation: bool, default=True
+        If `True`, we obtain Rop via double application of Lop.
+        If `False`, the legacy Rop implementation is used. The number of graphs that support this form
+        is much more restricted, and the generated graphs may be less optimized.
+
+    Returns
+    -------
+    :class:`~pytensor.graph.basic.Variable` or list/tuple of Variables
+        A symbolic expression such obeying
+        ``R_op[i] = sum_j (d f[i] / d wrt[j]) eval_point[j]``,
+        where the indices in that expression are magic multidimensional
+        indices that specify both the position within a list and all
+        coordinates of the tensor elements.
+        If `f` is a list/tuple, then return a list/tuple with the results.
+
+    References
+    ----------
+    .. [1] J. Towns, "A new trick for calculating Jacobian vector products", 2017.
+           Available: https://j-towns.github.io/2017/06/12/A-new-trick.html
+    """
+
+    if not isinstance(wrt, list | tuple):
+        _wrt: list[Variable] = [pytensor.tensor.as_tensor_variable(wrt)]
+    else:
+        _wrt = [pytensor.tensor.as_tensor_variable(x) for x in wrt]
+
+    if not isinstance(eval_points, list | tuple):
+        _eval_points: list[Variable] = [pytensor.tensor.as_tensor_variable(eval_points)]
+    else:
+        _eval_points = [pytensor.tensor.as_tensor_variable(x) for x in eval_points]
+
+    if not isinstance(f, list | tuple):
+        _f: list[Variable] = [pytensor.tensor.as_tensor_variable(f)]
+    else:
+        _f = [pytensor.tensor.as_tensor_variable(x) for x in f]
+
+    if len(_wrt) != len(_eval_points):
+        raise ValueError("`wrt` must be the same length as `eval_points`.")
+
+    # Check that each element of wrt corresponds to an element
+    # of eval_points with the same dimensionality.
+    for i, (wrt_elem, eval_point) in enumerate(zip(_wrt, _eval_points, strict=True)):
+        try:
+            if wrt_elem.type.ndim != eval_point.type.ndim:
+                raise ValueError(
+                    f"Elements {i} of `wrt` and `eval_point` have mismatched dimensionalities: "
+                    f"{wrt_elem.type.ndim} and {eval_point.type.ndim}"
+                )
+        except AttributeError:
+            # wrt_elem and eval_point don't always have ndim like random type
+            # Tensor, Sparse have the ndim attribute
+            pass
+
+    if use_op_rop_implementation:
+        rval = _rop_legacy(
+            _f, _wrt, _eval_points, disconnected_outputs, return_disconnected
+        )
+    else:
+        rval = pushforward_through_pullback(
+            _f, _wrt, _eval_points, disconnected_outputs, return_disconnected
+        )
+
     using_list = isinstance(f, list)
     using_tuple = isinstance(f, tuple)
     return as_list_or_tuple(using_list, using_tuple, rval)
@@ -348,6 +463,7 @@ def Lop(
     eval_points: Variable | Sequence[Variable],
     consider_constant: Sequence[Variable] | None = None,
     disconnected_inputs: Literal["ignore", "warn", "raise"] = "raise",
+    return_disconnected: Literal["none", "zero", "disconnected"] = "zero",
 ) -> Variable | None | Sequence[Variable | None]:
     """Computes the L-operator applied to `f` with respect to `wrt` at `eval_points`.
 
@@ -404,6 +520,7 @@ def Lop(
         consider_constant=consider_constant,
         wrt=_wrt,
         disconnected_inputs=disconnected_inputs,
+        return_disconnected=return_disconnected,
     )
 
     using_list = isinstance(wrt, list)
diff --git a/pytensor/scan/op.py b/pytensor/scan/op.py
index a01347ef9..1dbc93b9f 100644
--- a/pytensor/scan/op.py
+++ b/pytensor/scan/op.py
@@ -72,6 +72,7 @@ from pytensor.graph.basic import (
 from pytensor.graph.features import NoOutputFromInplace
 from pytensor.graph.op import HasInnerGraph, Op
 from pytensor.graph.replace import clone_replace
+from pytensor.graph.type import HasShape
 from pytensor.graph.utils import InconsistencyError, MissingInputError
 from pytensor.link.c.basic import CLinker
 from pytensor.printing import op_debug_information
@@ -2509,13 +2510,25 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             return rval
 
         var_mappings = self.get_oinp_iinp_iout_oout_mappings()
-        dC_dinps_t = [None for inp in diff_inputs]
         disconnected_dC_dinps_t = [True for inp in diff_inputs]
+
+        n_mit_mot_outs = info.n_mit_mot_outs
+        # In the case of mit-mot there can be more inner outputs than outer ones
+        n_extra_mit_mot_outs = n_mit_mot_outs - info.n_mit_mot
+        idx_nitsot_out_start = n_mit_mot_outs + info.n_mit_sot + info.n_sit_sot
+        idx_nitsot_out_end = idx_nitsot_out_start + info.n_nit_sot
+
+        # Create dummy variables for the internal input gradients
+        states = (
+            self.inner_mitmot(self_inputs)
+            + self.inner_mitsot(self_inputs)
+            + self.inner_sitsot(self_inputs)
+        )
         dC_dXts = []
         Xts = []
         for idx, Xt in enumerate(diff_outputs):
             # We are looking for x[t-1] for a given x[t]
-            if idx >= info.n_mit_mot_outs:
+            if idx >= n_mit_mot_outs:
                 Xt_placeholder = safe_new(Xt)
                 Xts.append(Xt_placeholder)
 
@@ -2523,9 +2536,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             # or not. NOTE : This cannot be done by using
             # "if Xt not in self.inner_nitsot_outs(self_outputs)" because
             # the exact same variable can be used as multiple outputs.
-            idx_nitsot_start = info.n_mit_mot + info.n_mit_sot + info.n_sit_sot
-            idx_nitsot_end = idx_nitsot_start + info.n_nit_sot
-            if idx < idx_nitsot_start or idx >= idx_nitsot_end:
+            if idx < idx_nitsot_out_start or idx >= idx_nitsot_out_end:
                 # What we do here is loop through dC_douts and collect all
                 # those that are connected to the specific one and do an
                 # upcast on all of their dtypes to get the dtype for this
@@ -2533,12 +2544,6 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                 # specific previous step is defined or not is done somewhere
                 # else.
                 dtypes = []
-                states = (
-                    self.inner_mitmot(self_inputs)
-                    + self.inner_mitsot(self_inputs)
-                    + self.inner_sitsot(self_inputs)
-                )
-
                 for pos, inp in enumerate(states):
                     if inp in graph_inputs([Xt]):
                         # Get the index of the outer output that to which
@@ -2555,35 +2560,43 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                     new_dtype = config.floatX
                 dC_dXt = safe_new(Xt, dtype=new_dtype)
             else:
-                if isinstance(dC_douts[idx].type, DisconnectedType):
+                # nit-sot outputs
+                # If not disconnected assume the output gradient type is a valid type for the input gradient
+                if isinstance(
+                    dC_douts[idx - n_extra_mit_mot_outs].type, DisconnectedType
+                ):
                     continue
-                dC_dXt = safe_new(dC_douts[idx][0])
+                dC_dXt = safe_new(dC_douts[idx - n_extra_mit_mot_outs][0])
             dC_dXts.append(dC_dXt)
 
+        # Handle cases where the very same variable may be used as different outputs
+        # TODO: Couldn't we add a view Op to avoid this when building the Scan graph?
         known_grads = {}
         dc_dxts_idx = 0
         for i in range(len(diff_outputs)):
-            if i < idx_nitsot_start or i >= idx_nitsot_end:
-                if diff_outputs[i] in known_grads:
-                    known_grads[diff_outputs[i]] += dC_dXts[dc_dxts_idx]
-                else:
-                    known_grads[diff_outputs[i]] = dC_dXts[dc_dxts_idx]
-                dc_dxts_idx += 1
+            if not (i < idx_nitsot_out_start or i >= idx_nitsot_out_end) and isinstance(
+                dC_douts[i - n_extra_mit_mot_outs].type, DisconnectedType
+            ):
+                # Special case where we don't have a dC_dXt for disconnected nitsot outputs
+                continue
+
+            # Just some trouble to avoid a +0
+            if diff_outputs[i] in known_grads:
+                known_grads[diff_outputs[i]] += dC_dXts[dc_dxts_idx]
             else:
-                if isinstance(dC_douts[i].type, DisconnectedType):
-                    continue
-                else:
-                    if diff_outputs[i] in known_grads:
-                        known_grads[diff_outputs[i]] += dC_dXts[dc_dxts_idx]
-                    else:
-                        known_grads[diff_outputs[i]] = dC_dXts[dc_dxts_idx]
-                    dc_dxts_idx += 1
+                known_grads[diff_outputs[i]] = dC_dXts[dc_dxts_idx]
+            dc_dxts_idx += 1
+
         dC_dinps_t = compute_all_gradients(known_grads)
 
         # mask inputs that get no gradients
         for dx in range(len(dC_dinps_t)):
-            if not dC_dinps_t[dx]:
-                dC_dinps_t[dx] = pt.zeros_like(diff_inputs[dx])
+            if dC_dinps_t[dx] is None:
+                dC_dinps_t[dx] = dC_dinps_t[dx] = (
+                    pt.zeros_like(diff_inputs[dx])
+                    if isinstance(diff_inputs[dx].type, HasShape)
+                    else pt.zeros(())
+                )
             else:
                 disconnected_dC_dinps_t[dx] = False
                 for Xt, Xt_placeholder in zip(
@@ -2846,7 +2859,6 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         for idx in range(info.n_sit_sot):
             mitmot_inp_taps.append([0, 1])
             mitmot_out_taps.append([1])
-            through_shared = False
             if not isinstance(dC_douts[idx + offset].type, DisconnectedType):
                 outer_inp_mitmot.append(dC_douts[idx + offset][::-1])
             else:
@@ -2958,7 +2970,8 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             else:
                 outer_inp_sitsot.append(
                     pt.zeros(
-                        [grad_steps + 1] + [x.shape[i] for i in range(x.ndim)],
+                        [grad_steps + 1]
+                        + (list(x.shape) if isinstance(x.type, HasShape) else []),
                         dtype=y.dtype,
                     )
                 )
@@ -3007,9 +3020,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             name=f"grad_of_{self.name}" if self.name else None,
             allow_gc=self.allow_gc,
         )
-        outputs = local_op(*outer_inputs)
-        if not isinstance(outputs, list | tuple):
-            outputs = [outputs]
+        outputs = local_op(*outer_inputs, return_list=True)
         # Re-order the gradients correctly
         gradients = [DisconnectedType()()]
 
@@ -3095,7 +3106,6 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                     )
                 )
 
-        start = len(gradients)
         gradients += [DisconnectedType()() for _ in range(info.n_nit_sot)]
         begin = end
 
@@ -3155,7 +3165,12 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             rop_self_outputs = self_outputs
         if info.n_shared_outs > 0:
             rop_self_outputs = rop_self_outputs[: -info.n_shared_outs]
-        rop_outs = Rop(rop_self_outputs, rop_of_inputs, inner_eval_points)
+        rop_outs = Rop(
+            rop_self_outputs,
+            rop_of_inputs,
+            inner_eval_points,
+            use_op_rop_implementation=True,
+        )
         if not isinstance(rop_outs, list | tuple):
             rop_outs = [rop_outs]
         # Step 2. Figure out what corresponds to what in the scan
diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py
index b185f686b..9fa823feb 100644
--- a/pytensor/tensor/math.py
+++ b/pytensor/tensor/math.py
@@ -431,20 +431,25 @@ class Max(NonZeroDimsCAReduce):
         return (g_x,)
 
     def R_op(self, inputs, eval_points):
+        [x] = inputs
         if eval_points[0] is None:
-            return [None, None]
-        if len(self.axis) != 1:
-            raise ValueError("R_op supported for max only for one axis!")
-        if self.axis[0] > 1:
-            raise ValueError("R_op supported for max only when axis is 0 or 1")
+            return [None]
+        axis = tuple(range(x.ndim) if self.axis is None else self.axis)
+        if isinstance(axis, int):
+            axis = [axis]
+        if len(axis) != 1:
+            raise NotImplementedError("R_op supported for max only for one axis!")
+        if axis[0] > 1:
+            raise NotImplementedError("R_op supported for max only when axis is 0 or 1")
         if inputs[0].ndim != 2:
-            raise ValueError("R_op supported for max only when input is a matrix")
-        max_pos = Argmax(self.axis).make_node(*inputs).outputs
-        # print(eval_points[0].eval())
+            raise NotImplementedError(
+                "R_op supported for max only when input is a matrix"
+            )
+        max_pos = Argmax(self.axis)(*inputs)
         if self.axis[0] == 0:
-            return [eval_points[0][max_pos, arange(eval_points[0].shape[1])], None]
+            return [eval_points[0][max_pos, arange(eval_points[0].shape[1])]]
         else:
-            return [eval_points[0][arange(eval_points[0].shape[0]), max_pos], None]
+            return [eval_points[0][arange(eval_points[0].shape[0]), max_pos]]
 
 
 class Min(NonZeroDimsCAReduce):


Test patch:
diff --git a/tests/compile/test_builders.py b/tests/compile/test_builders.py
index 8fc2a529d..ba0257cdd 100644
--- a/tests/compile/test_builders.py
+++ b/tests/compile/test_builders.py
@@ -306,7 +306,8 @@ class TestOpFromGraph(unittest_tools.InferShapeTester):
     @pytest.mark.parametrize(
         "cls_ofg", [OpFromGraph, partial(OpFromGraph, inline=True)]
     )
-    def test_rop(self, cls_ofg):
+    @pytest.mark.parametrize("use_op_rop_implementation", [True, False])
+    def test_rop(self, cls_ofg, use_op_rop_implementation):
         a = vector()
         M = matrix()
         b = dot(a, M)
@@ -315,7 +316,7 @@ class TestOpFromGraph(unittest_tools.InferShapeTester):
         W = matrix()
         y = op_matmul(x, W)
         du = vector()
-        dv = Rop(y, x, du)
+        dv = Rop(y, x, du, use_op_rop_implementation=use_op_rop_implementation)
         fn = function([x, W, du], dv)
         xval = np.random.random((16,)).astype(config.floatX)
         Wval = np.random.random((16, 16)).astype(config.floatX)
@@ -324,7 +325,8 @@ class TestOpFromGraph(unittest_tools.InferShapeTester):
         dvval2 = fn(xval, Wval, duval)
         np.testing.assert_array_almost_equal(dvval2, dvval, 4)
 
-    def test_rop_multiple_outputs(self):
+    @pytest.mark.parametrize("use_op_rop_implementation", [True, False])
+    def test_rop_multiple_outputs(self, use_op_rop_implementation):
         a = vector()
         M = matrix()
         b = dot(a, M)
@@ -339,21 +341,21 @@ class TestOpFromGraph(unittest_tools.InferShapeTester):
         duval = np.random.random((16,)).astype(config.floatX)
 
         y = op_matmul(x, W)[0]
-        dv = Rop(y, x, du)
+        dv = Rop(y, x, du, use_op_rop_implementation=use_op_rop_implementation)
         fn = function([x, W, du], dv)
         result_dvval = fn(xval, Wval, duval)
         expected_dvval = np.dot(duval, Wval)
         np.testing.assert_array_almost_equal(result_dvval, expected_dvval, 4)
 
         y = op_matmul(x, W)[1]
-        dv = Rop(y, x, du)
+        dv = Rop(y, x, du, use_op_rop_implementation=use_op_rop_implementation)
         fn = function([x, W, du], dv)
         result_dvval = fn(xval, Wval, duval)
         expected_dvval = -np.dot(duval, Wval)
         np.testing.assert_array_almost_equal(result_dvval, expected_dvval, 4)
 
         y = pt.add(*op_matmul(x, W))
-        dv = Rop(y, x, du)
+        dv = Rop(y, x, du, use_op_rop_implementation=use_op_rop_implementation)
         fn = function([x, W, du], dv)
         result_dvval = fn(xval, Wval, duval)
         expected_dvval = np.zeros_like(np.dot(duval, Wval))
@@ -362,7 +364,16 @@ class TestOpFromGraph(unittest_tools.InferShapeTester):
     @pytest.mark.parametrize(
         "cls_ofg", [OpFromGraph, partial(OpFromGraph, inline=True)]
     )
-    def test_rop_override(self, cls_ofg):
+    @pytest.mark.parametrize(
+        "use_op_rop_implementation",
+        [
+            True,
+            pytest.param(
+                False, marks=pytest.mark.xfail(reason="Custom ROp is ignored")
+            ),
+        ],
+    )
+    def test_rop_override(self, cls_ofg, use_op_rop_implementation):
         x, y = vectors("xy")
 
         def ro(inps, epts):
@@ -380,7 +391,12 @@ class TestOpFromGraph(unittest_tools.InferShapeTester):
         du, dv = vector("du"), vector("dv")
         for op in [op_mul, op_mul2]:
             zz = op_mul(xx, yy)
-            dw = Rop(zz, [xx, yy], [du, dv])
+            dw = Rop(
+                zz,
+                [xx, yy],
+                [du, dv],
+                use_op_rop_implementation=use_op_rop_implementation,
+            )
             fn = function([xx, yy, du, dv], dw)
             vals = np.random.random((4, 32)).astype(config.floatX)
             dwval = fn(*vals)
diff --git a/tests/scan/test_basic.py b/tests/scan/test_basic.py
index 9fa893ab2..351c2e703 100644
--- a/tests/scan/test_basic.py
+++ b/tests/scan/test_basic.py
@@ -1922,7 +1922,8 @@ class TestScan:
         fgrad = function([], g_sh)
         assert fgrad() == 1
 
-    def test_R_op(self):
+    @pytest.mark.parametrize("use_op_rop_implementation", [True, False])
+    def test_R_op(self, use_op_rop_implementation):
         seed = utt.fetch_seed()
         rng = np.random.default_rng(seed)
         floatX = config.floatX
@@ -1957,9 +1958,9 @@ class TestScan:
         eh0 = vector("eh0")
         eW = matrix("eW")
 
-        nwo_u = Rop(o, _u, eu)
-        nwo_h0 = Rop(o, _h0, eh0)
-        nwo_W = Rop(o, _W, eW)
+        nwo_u = Rop(o, _u, eu, use_op_rop_implementation=use_op_rop_implementation)
+        nwo_h0 = Rop(o, _h0, eh0, use_op_rop_implementation=use_op_rop_implementation)
+        nwo_W = Rop(o, _W, eW, use_op_rop_implementation=use_op_rop_implementation)
         fn_rop = function(
             [u, h0, W, eu, eh0, eW], [nwo_u, nwo_h0, nwo_W], on_unused_input="ignore"
         )
@@ -1992,12 +1993,13 @@ class TestScan:
         vnu, vnh0, vnW = fn_rop(v_u, v_h0, v_W, v_eu, v_eh0, v_eW)
         tnu, tnh0, tnW = fn_test(v_u, v_h0, v_W, v_eu, v_eh0, v_eW)
 
-        utt.assert_allclose(vnu, tnu, atol=1e-6)
-        utt.assert_allclose(vnh0, tnh0, atol=1e-6)
-        utt.assert_allclose(vnW, tnW, atol=1e-6)
+        np.testing.assert_allclose(vnu, tnu, atol=1e-6)
+        np.testing.assert_allclose(vnh0, tnh0, atol=1e-6)
+        np.testing.assert_allclose(vnW, tnW, atol=1e-6)
 
     @pytest.mark.slow
-    def test_R_op_2(self):
+    @pytest.mark.parametrize("use_op_rop_implementation", [True, False])
+    def test_R_op_2(self, use_op_rop_implementation):
         seed = utt.fetch_seed()
         rng = np.random.default_rng(seed)
         floatX = config.floatX
@@ -2040,9 +2042,9 @@ class TestScan:
         eh0 = vector("eh0")
         eW = matrix("eW")
 
-        nwo_u = Rop(o, _u, eu)
-        nwo_h0 = Rop(o, _h0, eh0)
-        nwo_W = Rop(o, _W, eW)
+        nwo_u = Rop(o, _u, eu, use_op_rop_implementation=use_op_rop_implementation)
+        nwo_h0 = Rop(o, _h0, eh0, use_op_rop_implementation=use_op_rop_implementation)
+        nwo_W = Rop(o, _W, eW, use_op_rop_implementation=use_op_rop_implementation)
         fn_rop = function(
             [u, h0, W, eu, eh0, eW], [nwo_u, nwo_h0, nwo_W, o], on_unused_input="ignore"
         )
@@ -2074,11 +2076,12 @@ class TestScan:
         )
 
         tnu, tnh0, tnW, tno = fn_test(v_u, v_h0, v_W, v_eu, v_eh0, v_eW)
-        utt.assert_allclose(vnu, tnu, atol=1e-6)
-        utt.assert_allclose(vnh0, tnh0, atol=1e-6)
-        utt.assert_allclose(vnW, tnW, atol=2e-6)
+        np.testing.assert_allclose(vnu, tnu, atol=1e-6)
+        np.testing.assert_allclose(vnh0, tnh0, atol=1e-6)
+        np.testing.assert_allclose(vnW, tnW, atol=2e-6)
 
-    def test_R_op_mitmot(self):
+    @pytest.mark.parametrize("use_op_rop_implementation", [True, False])
+    def test_R_op_mitmot(self, use_op_rop_implementation):
         # this test is a copy paste from the script given by Justin Bayer to
         # reproduce this bug
         # We have 2 parameter groups with the following shapes.
@@ -2094,13 +2097,10 @@ class TestScan:
         W1 = pars[:3].reshape(W1shape)
         W2 = pars[3:].reshape(W2shape)
 
-        # Define recurrent model. We are using a model where each input is a
-        # tensor
-        # of shape (T, B, D) where T is the number of timesteps, B is the
-        # number of
-        # sequences iterated over in parallel and D is the dimensionality of
-        # each
-        # item at a timestep.
+        # Define recurrent model. We are using a model where each input
+        # is a tensor of shape (T, B, D) where T is the number of timesteps,
+        # B is the number of sequences iterated over in parallel and
+        # D is the dimensionality of each item at a timestep.
 
         inpt = tensor3("inpt")
         target = tensor3("target")
@@ -2128,7 +2128,130 @@ class TestScan:
         d_cost_wrt_pars = grad(cost, pars)
 
         p = dvector()
-        Rop(d_cost_wrt_pars, pars, p)
+        # TODO: We should test something about the Rop!
+        Rop(
+            d_cost_wrt_pars,
+            pars,
+            p,
+            use_op_rop_implementation=use_op_rop_implementation,
+        )
+
+    def test_second_derivative_disconnected_cost_with_mit_mot(self):
+        # This test is a regression test for a bug that was revealed
+        # when we computed the pushforward of a Scan gradient via two applications of pullback
+        seq = pt.vector("seq", shape=(2,))
+        z = pt.scalar("z")
+        x0 = pt.vector("x0", shape=(2,))
+
+        # When s is 1 and z is 2, xs[-1] is just a sneaky
+        # x ** 4 (after two nsteps)
+        # grad should be 4 * x ** 3
+        # and grad of grad should be 12 * x ** 2
+        def step(s, xtm2, xtm1, z):
+            return s * ((xtm2 * 0 + xtm1) ** 2) * (z / 2)
+
+        xs, _ = scan(
+            step,
+            sequences=[seq],
+            outputs_info=[{"initial": x0, "taps": (-2, -1)}],
+            non_sequences=[z],
+            n_steps=2,
+        )
+        last_x = xs[-1]
+
+        g_wrt_x0, g_wrt_z, g_wrt_seq = pt.grad(last_x, [x0, z, seq])
+        g = g_wrt_x0.sum() + g_wrt_z.sum() * 0 + g_wrt_seq.sum() * 0
+        assert g.eval({seq: [1, 1], x0: [1, 1], z: 2}) == 4
+        gg = pt.grad(g, wrt=x0).sum()
+        assert gg.eval({seq: [1, 1], x0: [1, 1], z: 2}) == 12
+        assert gg.eval({seq: [2, 2], x0: [1, 1], z: 2}) == 96
+
+        # Leave out z
+        g_wrt_x0, g_wrt_seq = pt.grad(last_x, [x0, seq])
+        g = g_wrt_x0.sum() + g_wrt_seq.sum() * 0
+        gg = pt.grad(g, wrt=x0).sum()
+        assert gg.eval({seq: [1, 1], x0: [1, 1], z: 2}) == 12
+        assert gg.eval({seq: [2, 2], x0: [1, 1], z: 2}) == 96
+
+        # Leave out seq
+        g_wrt_x0, g_wrt_z = pt.grad(last_x, [x0, z])
+        g = g_wrt_x0.sum() + g_wrt_z.sum() * 0
+        gg = pt.grad(g, wrt=x0).sum()
+        assert gg.eval({seq: [1, 1], x0: [1, 1], z: 2}) == 12
+        assert gg.eval({seq: [1, 1], x0: [1, 1], z: 1}) == 3 / 2
+
+        # Leave out z and seq
+        g_wrt_x0 = pt.grad(last_x, x0)
+        g = g_wrt_x0.sum()
+        gg = pt.grad(g, wrt=x0).sum()
+        assert gg.eval({seq: [1, 1], x0: [1, 1], z: 2}) == 12
+        assert gg.eval({seq: [1, 1], x0: [1, 1], z: 1}) == 3 / 2
+
+    @pytest.mark.parametrize("case", ("inside-explicit", "inside-implicit", "outside"))
+    def test_non_shaped_input_disconnected_gradient(self, case):
+        """Test that Scan gradient works when non shaped variables are disconnected from the gradient.
+
+        Regression test for https://github.com/pymc-devs/pytensor/issues/6
+        """
+
+        # In all cases rng is disconnected from the output gradient
+        # Note that when it is an input to the scan (explicit or not) it is still not updated by the scan,
+        # so it is equivalent to the `outside` case. A rewrite could have legally hoisted the rng out of the scan.
+        rng = shared(np.random.default_rng())
+
+        data = pt.zeros(16)
+
+        nonlocal_random_index = pt.random.integers(16, rng=rng)
+        nonlocal_random_datum = data[nonlocal_random_index]
+
+        if case == "outside":
+
+            def step(s, random_datum):
+                return (random_datum + s) ** 2
+
+            strict = True
+            non_sequences = [nonlocal_random_datum]
+
+        elif case == "inside-implicit":
+
+            def step(s):
+                return (nonlocal_random_datum + s) ** 2
+
+            strict = False
+            non_sequences = []  # Scan will introduce the non_sequences for us
+
+        elif case == "inside-explicit":
+
+            def step(s, data, rng):
+                random_index = pt.random.integers(
+                    16, rng=rng
+                )  # Not updated by the scan
+                random_datum = data[random_index]
+                return (random_datum + s) ** 2
+
+            strict = (True,)
+            non_sequences = [data, rng]
+
+        else:
+            raise ValueError(f"Invalid case: {case}")
+
+        seq = vector("seq")
+        xs, _ = scan(
+            step,
+            sequences=[seq],
+            non_sequences=non_sequences,
+            strict=strict,
+        )
+        x0 = xs[0]
+
+        np.testing.assert_allclose(
+            x0.eval({seq: [np.pi, np.nan, np.nan]}),
+            np.pi**2,
+        )
+        np.testing.assert_allclose(
+            grad(x0, seq)[0].eval({seq: [np.pi, np.nan, np.nan]}),
+            2 * np.pi,
+        )
 
 
 @pytest.mark.skipif(
diff --git a/tests/tensor/rewriting/test_linalg.py b/tests/tensor/rewriting/test_linalg.py
index c9b9afff1..50e48ce95 100644
--- a/tests/tensor/rewriting/test_linalg.py
+++ b/tests/tensor/rewriting/test_linalg.py
@@ -14,7 +14,7 @@ from pytensor.graph.rewriting.utils import rewrite_graph
 from pytensor.tensor import swapaxes
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
-from pytensor.tensor.math import _allclose, dot, matmul
+from pytensor.tensor.math import dot, matmul
 from pytensor.tensor.nlinalg import (
     SVD,
     Det,
@@ -42,18 +42,19 @@ from tests import unittest_tools as utt
 from tests.test_rop import break_op
 
 
-ATOL = RTOL = 1e-3 if config.floatX == "float32" else 1e-8
-
-
-def test_rop_lop():
+def test_matrix_inverse_rop_lop():
+    rtol = 1e-7 if config.floatX == "float64" else 1e-5
     mx = matrix("mx")
     mv = matrix("mv")
     v = vector("v")
     y = MatrixInverse()(mx).sum(axis=0)
 
-    yv = pytensor.gradient.Rop(y, mx, mv)
+    yv = pytensor.gradient.Rop(y, mx, mv, use_op_rop_implementation=True)
     rop_f = function([mx, mv], yv)
 
+    yv_via_lop = pytensor.gradient.Rop(y, mx, mv, use_op_rop_implementation=False)
+    rop_via_lop_f = function([mx, mv], yv_via_lop)
+
     sy, _ = pytensor.scan(
         lambda i, y, x, v: (pytensor.gradient.grad(y[i], x) * v).sum(),
         sequences=pt.arange(y.shape[0]),
@@ -65,22 +66,16 @@ def test_rop_lop():
     vx = np.asarray(rng.standard_normal((4, 4)), pytensor.config.floatX)
     vv = np.asarray(rng.standard_normal((4, 4)), pytensor.config.floatX)
 
-    v1 = rop_f(vx, vv)
-    v2 = scan_f(vx, vv)
+    v_ref = scan_f(vx, vv)
+    np.testing.assert_allclose(rop_f(vx, vv), v_ref, rtol=rtol)
+    np.testing.assert_allclose(rop_via_lop_f(vx, vv), v_ref, rtol=rtol)
 
-    assert _allclose(v1, v2), f"ROP mismatch: {v1} {v2}"
-
-    raised = False
-    try:
+    with pytest.raises(ValueError):
         pytensor.gradient.Rop(
-            pytensor.clone_replace(y, replace={mx: break_op(mx)}), mx, mv
-        )
-    except ValueError:
-        raised = True
-    if not raised:
-        raise Exception(
-            "Op did not raised an error even though the function"
-            " is not differentiable"
+            pytensor.clone_replace(y, replace={mx: break_op(mx)}),
+            mx,
+            mv,
+            use_op_rop_implementation=True,
         )
 
     vv = np.asarray(rng.uniform(size=(4,)), pytensor.config.floatX)
@@ -90,9 +85,9 @@ def test_rop_lop():
     sy = pytensor.gradient.grad((v * y).sum(), mx)
     scan_f = function([mx, v], sy)
 
-    v1 = lop_f(vx, vv)
-    v2 = scan_f(vx, vv)
-    assert _allclose(v1, v2), f"LOP mismatch: {v1} {v2}"
+    v_ref = scan_f(vx, vv)
+    v = lop_f(vx, vv)
+    np.testing.assert_allclose(v, v_ref, rtol=rtol)
 
 
 def test_transinv_to_invtrans():
@@ -630,11 +625,12 @@ def test_inv_diag_from_eye_mul(shape, inv_op):
     inverse_matrix = np.linalg.inv(x_test_matrix)
     rewritten_inverse = f_rewritten(x_test)
 
+    atol = rtol = 1e-3 if config.floatX == "float32" else 1e-8
     assert_allclose(
         inverse_matrix,
         rewritten_inverse,
-        atol=ATOL,
-        rtol=RTOL,
+        atol=atol,
+        rtol=rtol,
     )
 
 
@@ -657,11 +653,12 @@ def test_inv_diag_from_diag(inv_op):
     inverse_matrix = np.linalg.inv(x_test_matrix)
     rewritten_inverse = f_rewritten(x_test)
 
+    atol = rtol = 1e-3 if config.floatX == "float32" else 1e-8
     assert_allclose(
         inverse_matrix,
         rewritten_inverse,
-        atol=ATOL,
-        rtol=RTOL,
+        atol=atol,
+        rtol=rtol,
     )
 
 
diff --git a/tests/tensor/test_shape.py b/tests/tensor/test_shape.py
index e85b8cfd4..7700d2b14 100644
--- a/tests/tensor/test_shape.py
+++ b/tests/tensor/test_shape.py
@@ -603,7 +603,7 @@ class TestSpecifyBroadcastable:
 
 class TestRopLop(RopLopChecker):
     def test_shape(self):
-        self.check_nondiff_rop(self.x.shape[0])
+        self.check_nondiff_rop(self.x.shape[0], self.x, self.v)
 
     def test_specifyshape(self):
         self.check_rop_lop(specify_shape(self.x, self.in_shape), self.in_shape)
diff --git a/tests/test_rop.py b/tests/test_rop.py
index 0b9fe41a1..b592f557a 100644
--- a/tests/test_rop.py
+++ b/tests/test_rop.py
@@ -16,8 +16,14 @@ import pytest
 
 import pytensor
 import pytensor.tensor as pt
-from pytensor import function
-from pytensor.gradient import Lop, Rop, grad, grad_undefined
+from pytensor import config, function
+from pytensor.gradient import (
+    Lop,
+    NullTypeGradError,
+    Rop,
+    grad,
+    grad_undefined,
+)
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
 from pytensor.tensor.math import argmax, dot
@@ -61,6 +67,10 @@ class RopLopChecker:
     Rop to class that inherit from it.
     """
 
+    @staticmethod
+    def rtol():
+        return 1e-7 if config.floatX == "float64" else 1e-5
+
     def setup_method(self):
         # Using vectors make things a lot simpler for generating the same
         # computations using scan
@@ -72,13 +82,13 @@ class RopLopChecker:
         self.mv = matrix("mv")
         self.mat_in_shape = (5 + self.rng.integers(3), 5 + self.rng.integers(3))
 
-    def check_nondiff_rop(self, y):
+    def check_nondiff_rop(self, y, x, v):
         """
         If your op is not differentiable(so you can't define Rop)
         test that an error is raised.
         """
         with pytest.raises(ValueError):
-            Rop(y, self.x, self.v)
+            Rop(y, x, v, use_op_rop_implementation=True)
 
     def check_mat_rop_lop(self, y, out_shape):
         """
@@ -106,8 +116,14 @@ class RopLopChecker:
         vv = np.asarray(
             self.rng.uniform(size=self.mat_in_shape), pytensor.config.floatX
         )
-        yv = Rop(y, self.mx, self.mv)
+        yv = Rop(y, self.mx, self.mv, use_op_rop_implementation=True)
         rop_f = function([self.mx, self.mv], yv, on_unused_input="ignore")
+
+        yv_through_lop = Rop(y, self.mx, self.mv, use_op_rop_implementation=False)
+        rop_through_lop_f = function(
+            [self.mx, self.mv], yv_through_lop, on_unused_input="ignore"
+        )
+
         sy, _ = pytensor.scan(
             lambda i, y, x, v: (grad(y[i], x) * v).sum(),
             sequences=pt.arange(y.shape[0]),
@@ -115,13 +131,14 @@ class RopLopChecker:
         )
         scan_f = function([self.mx, self.mv], sy, on_unused_input="ignore")
 
-        v1 = rop_f(vx, vv)
-        v2 = scan_f(vx, vv)
-
-        assert np.allclose(v1, v2), f"ROP mismatch: {v1} {v2}"
+        v_ref = scan_f(vx, vv)
+        np.testing.assert_allclose(rop_f(vx, vv), v_ref)
+        np.testing.assert_allclose(rop_through_lop_f(vx, vv), v_ref)
 
         self.check_nondiff_rop(
-            pytensor.clone_replace(y, replace={self.mx: break_op(self.mx)})
+            pytensor.clone_replace(y, replace={self.mx: break_op(self.mx)}),
+            self.mx,
+            self.mv,
         )
 
         vv = np.asarray(self.rng.uniform(size=out_shape), pytensor.config.floatX)
@@ -131,45 +148,47 @@ class RopLopChecker:
         sy = grad((self.v * y).sum(), self.mx)
         scan_f = function([self.mx, self.v], sy)
 
-        v1 = lop_f(vx, vv)
-        v2 = scan_f(vx, vv)
-        assert np.allclose(v1, v2), f"LOP mismatch: {v1} {v2}"
+        v = lop_f(vx, vv)
+        v_ref = scan_f(vx, vv)
+        np.testing.assert_allclose(v, v_ref)
 
-    def check_rop_lop(self, y, out_shape):
+    def check_rop_lop(self, y, out_shape, check_nondiff_rop: bool = True):
         """
         As check_mat_rop_lop, except the input is self.x which is a
         vector. The output is still a vector.
         """
+        rtol = self.rtol()
+
         # TEST ROP
         vx = np.asarray(self.rng.uniform(size=self.in_shape), pytensor.config.floatX)
         vv = np.asarray(self.rng.uniform(size=self.in_shape), pytensor.config.floatX)
 
-        yv = Rop(y, self.x, self.v)
+        yv = Rop(y, self.x, self.v, use_op_rop_implementation=True)
         rop_f = function([self.x, self.v], yv, on_unused_input="ignore")
+
+        yv_through_lop = Rop(y, self.x, self.v, use_op_rop_implementation=False)
+        rop_through_lop_f = function(
+            [self.x, self.v], yv_through_lop, on_unused_input="ignore"
+        )
+
         J, _ = pytensor.scan(
             lambda i, y, x: grad(y[i], x),
             sequences=pt.arange(y.shape[0]),
             non_sequences=[y, self.x],
         )
         sy = dot(J, self.v)
-
         scan_f = function([self.x, self.v], sy, on_unused_input="ignore")
 
-        v1 = rop_f(vx, vv)
-        v2 = scan_f(vx, vv)
-        assert np.allclose(v1, v2), f"ROP mismatch: {v1} {v2}"
+        v_ref = scan_f(vx, vv)
+        np.testing.assert_allclose(rop_f(vx, vv), v_ref, rtol=rtol)
+        np.testing.assert_allclose(rop_through_lop_f(vx, vv), v_ref, rtol=rtol)
 
-        try:
-            Rop(
+        if check_nondiff_rop:
+            self.check_nondiff_rop(
                 pytensor.clone_replace(y, replace={self.x: break_op(self.x)}),
                 self.x,
                 self.v,
             )
-        except ValueError:
-            pytest.skip(
-                "Rop does not handle non-differentiable inputs "
-                "correctly. Bug exposed by fixing Add.grad method."
-            )
 
         vx = np.asarray(self.rng.uniform(size=self.in_shape), pytensor.config.floatX)
         vv = np.asarray(self.rng.uniform(size=out_shape), pytensor.config.floatX)
@@ -182,22 +201,20 @@ class RopLopChecker:
             non_sequences=[y, self.x],
         )
         sy = dot(self.v, J)
-
         scan_f = function([self.x, self.v], sy)
 
-        v1 = lop_f(vx, vv)
-        v2 = scan_f(vx, vv)
-        assert np.allclose(v1, v2), f"LOP mismatch: {v1} {v2}"
+        v = lop_f(vx, vv)
+        v_ref = scan_f(vx, vv)
+        np.testing.assert_allclose(v, v_ref, rtol=rtol)
 
 
 class TestRopLop(RopLopChecker):
     def test_max(self):
-        # self.check_mat_rop_lop(pt_max(self.mx, axis=[0,1])[0], ())
         self.check_mat_rop_lop(pt_max(self.mx, axis=0), (self.mat_in_shape[1],))
         self.check_mat_rop_lop(pt_max(self.mx, axis=1), (self.mat_in_shape[0],))
 
     def test_argmax(self):
-        self.check_nondiff_rop(argmax(self.mx, axis=1))
+        self.check_nondiff_rop(argmax(self.mx, axis=1), self.mx, self.mv)
 
     def test_subtensor(self):
         self.check_rop_lop(self.x[:4], (4,))
@@ -252,10 +269,14 @@ class TestRopLop(RopLopChecker):
         insh = self.in_shape[0]
         vW = np.asarray(self.rng.uniform(size=(insh, insh)), pytensor.config.floatX)
         W = pytensor.shared(vW)
-        self.check_rop_lop(dot(self.x, W), self.in_shape)
+        # check_nondiff_rop reveals an error in how legacy Rop handles non-differentiable paths
+        # See: test_Rop_partially_differentiable_paths
+        self.check_rop_lop(dot(self.x, W), self.in_shape, check_nondiff_rop=False)
 
     def test_elemwise0(self):
-        self.check_rop_lop((self.x + 1) ** 2, self.in_shape)
+        # check_nondiff_rop reveals an error in how legacy Rop handles non-differentiable paths
+        # See: test_Rop_partially_differentiable_paths
+        self.check_rop_lop((self.x + 1) ** 2, self.in_shape, check_nondiff_rop=False)
 
     def test_elemwise1(self):
         self.check_rop_lop(self.x + pt.cast(self.x, "int32"), self.in_shape)
@@ -287,18 +308,18 @@ class TestRopLop(RopLopChecker):
             self.mat_in_shape[0] * self.mat_in_shape[1] * self.in_shape[0],
         )
 
-    def test_invalid_input(self):
-        success = False
-
-        try:
-            Rop(0.0, [matrix()], [vector()])
-            success = True
-        except ValueError:
-            pass
-
-        assert not success
+    @pytest.mark.parametrize("use_op_rop_implementation", [True, False])
+    def test_invalid_input(self, use_op_rop_implementation):
+        with pytest.raises(ValueError):
+            Rop(
+                0.0,
+                [matrix()],
+                [vector()],
+                use_op_rop_implementation=use_op_rop_implementation,
+            )
 
-    def test_multiple_outputs(self):
+    @pytest.mark.parametrize("use_op_rop_implementation", [True, False])
+    def test_multiple_outputs(self, use_op_rop_implementation):
         m = matrix("m")
         v = vector("v")
         m_ = matrix("m_")
@@ -309,10 +330,20 @@ class TestRopLop(RopLopChecker):
         m_val = self.rng.uniform(size=(3, 7)).astype(pytensor.config.floatX)
         v_val = self.rng.uniform(size=(7,)).astype(pytensor.config.floatX)
 
-        rop_out1 = Rop([m, v, m + v], [m, v], [m_, v_])
+        rop_out1 = Rop(
+            [m, v, m + v],
+            [m, v],
+            [m_, v_],
+            use_op_rop_implementation=use_op_rop_implementation,
+        )
         assert isinstance(rop_out1, list)
         assert len(rop_out1) == 3
-        rop_out2 = Rop((m, v, m + v), [m, v], [m_, v_])
+        rop_out2 = Rop(
+            (m, v, m + v),
+            [m, v],
+            [m_, v_],
+            use_op_rop_implementation=use_op_rop_implementation,
+        )
         assert isinstance(rop_out2, tuple)
         assert len(rop_out2) == 3
 
@@ -322,12 +353,65 @@ class TestRopLop(RopLopChecker):
         f = pytensor.function([m, v, m_, v_], all_outs)
         f(mval, vval, m_val, v_val)
 
-    def test_Rop_dot_bug_18Oct2013_Jeremiah(self):
+    @pytest.mark.parametrize(
+        "use_op_rop_implementation",
+        [pytest.param(True, marks=pytest.mark.xfail()), False],
+    )
+    def test_Rop_partially_differentiable_paths(self, use_op_rop_implementation):
         # This test refers to a bug reported by Jeremiah Lowin on 18th Oct
         # 2013. The bug consists when through a dot operation there is only
         # one differentiable path (i.e. there is no gradient wrt to one of
         # the inputs).
         x = pt.arange(20.0).reshape([1, 20])
-        v = pytensor.shared(np.ones([20]))
+        v = pytensor.shared(np.ones([20]), name="v")
         d = dot(x, v).sum()
-        Rop(grad(d, v), v, v)
+
+        Rop(
+            grad(d, v),
+            v,
+            v,
+            use_op_rop_implementation=use_op_rop_implementation,
+            # 2025: This is a tricky case, the gradient of the gradient does not depend on v
+            # although v still exists in the graph inside a `Second` operator.
+            # The original test was checking that Rop wouldn't raise an error, but Lop does.
+            # Since the correct behavior is ambiguous, I let both implementations off the hook.
+            disconnected_outputs="raise" if use_op_rop_implementation else "ignore",
+        )
+
+        # 2025: Here is an unambiguous test for the original commented issue:
+        x = pt.matrix("x")
+        y = pt.matrix("y")
+        out = dot(x, break_op(y)).sum()
+        # Should not raise an error
+        Rop(
+            out,
+            [x],
+            [x.type()],
+            use_op_rop_implementation=use_op_rop_implementation,
+            disconnected_outputs="raise",
+        )
+
+        # More extensive testing shows that the legacy Rop implementation FAILS to raise when
+        # the cost is linked through strictly non-differentiable paths.
+        # This is not Dot specific, we would observe the same with any operation where the gradient
+        # with respect to one of the inputs does not depend on the original input (such as `mul`, `add`, ...)
+        out = dot(break_op(x), y).sum()
+        with pytest.raises((ValueError, NullTypeGradError)):
+            Rop(
+                out,
+                [x],
+                [x.type()],
+                use_op_rop_implementation=use_op_rop_implementation,
+                disconnected_outputs="raise",
+            )
+
+        # Only when both paths are non-differentiable is an error correctly raised again.
+        out = dot(break_op(x), break_op(y)).sum()
+        with pytest.raises((ValueError, NullTypeGradError)):
+            Rop(
+                out,
+                [x],
+                [x.type()],
+                use_op_rop_implementation=use_op_rop_implementation,
+                disconnected_outputs="raise",
+            )