Slow compilation time in warmup tests.

[zoj613]

Description of your problem or feature request

The warmup test, particularly

aehmc/tests/test_hmc.py

Line 49 in d54e2d0

def test_warmup_vector():

is very slow and this is due to the compilation time of warmup_fn.

Please provide a minimal, self-contained, and reproducible example.

python -m pytest tests/test_hmc.py::test_warmup_vector -v

Please provide the full traceback of any errors.
Running the test with profiling turned on produces the following:

Function profiling
==================
  Message:
  Time in 1 calls to Function.__call__: 1.529328e+01s
  Time in Function.vm.__call__: 15.29313341199304s (99.999%)
  Time in thunks: 15.29306435585022s (99.999%)
  Total compilation time: 1.663600e+02s
    Number of Apply nodes: 48
    Aesara rewrite time: 5.106596e+01s
       Aesara validate time: 3.606322e-02s
    Aesara Linker time (includes C, CUDA code generation/compiling): 115.27304236099008s
       Import time 2.332967e-01s
       Node make_thunk time 1.152699e+02s
           Node forall_inplace,cpu,window_adaptation}(TensorConstant{1000}, TensorConstant{[False Fal..lse  True]}, TensorConstant{[  0   1  ..7 998 999]}, Inc
Subtensor{InplaceSet;:int64:}.0, IncSubtensor{InplaceSet;:int64:}.0, IncSubtensor{InplaceSet;:int64:}.0, IncSubtensor{Set;:int64:}.0, IncSubtensor{Set;:int64:}
.0, IncSubtensor{Set;:int64:}.0, IncSubtensor{Set;:int64:}.0, IncSubtensor{InplaceSet;:int64:}.0, IncSubtensor{InplaceSet;:int64:}.0, IncSubtensor{InplaceSet;:
int64:}.0, IncSubtensor{InplaceSet;:int64:}.0, IncSubtensor{InplaceSet;:int64:}.0, IncSubtensor{InplaceSet;:int64:}.0, RandomGeneratorSharedVariable(<Generator
(PCG64) at 0x7F36B5424F20>), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F36B54265E0>), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F36B542
42E0>), RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F36B5427760>), IncSubtensor{InplaceSet;:int64:}.0, IncSubtensor{InplaceSet;:int64:}.0, IncSubtens
or{Set;:int64:}.0, IncSubtensor{InplaceSet;:int64:}.0) time 1.151257e+02s
           Node AllocEmpty{dtype='float64'}(TensorConstant{11}) time 1.173822e-02s
           Node InplaceDimShuffle{x}(Sum{acc_dtype=float64}.0) time 9.378754e-03s
           Node InplaceDimShuffle{x,0}(<TensorType(float64, (2,))>) time 9.352558e-03s
           Node Elemwise{Composite{(i0 - (i1 * i2))}}[(0, 2)](TensorConstant{(1,) of -2..2469692907}, TensorConstant{(1,) of 0.5}, InplaceDimShuffle{x}.0) time
 5.734075e-03s

Time in all call to aesara.grad() 1.157906e-01s
Time since aesara import 207.771s
Class
---
<% time> <sum %> <apply time> <time per call> <type> <#call> <#apply> <Class name>
  100.0%   100.0%      15.293s       1.53e+01s     Py       1       1   aesara.scan.op.Scan
   0.0%   100.0%       0.000s       1.33e-04s     Py       2       2   aesara.tensor.slinalg.SolveTriangular
   0.0%   100.0%       0.000s       2.24e-06s     C       17      17   aesara.tensor.subtensor.IncSubtensor
   0.0%   100.0%       0.000s       9.06e-06s     C        3       3   aesara.tensor.elemwise.DimShuffle
   0.0%   100.0%       0.000s       1.33e-06s     C       12      12   aesara.tensor.basic.AllocEmpty
   0.0%   100.0%       0.000s       6.20e-07s     C        5       5   aesara.tensor.subtensor.Subtensor
   0.0%   100.0%       0.000s       7.75e-07s     C        4       4   aesara.tensor.elemwise.Elemwise
   0.0%   100.0%       0.000s       1.91e-06s     C        1       1   aesara.tensor.math.Sum
   0.0%   100.0%       0.000s       3.18e-07s     C        3       3   aesara.tensor.shape.Unbroadcast
   ... (remaining 0 Classes account for   0.00%(0.00s) of the runtime)

Ops
---
<% time> <sum %> <apply time> <time per call> <type> <#call> <#apply> <Op name>
  100.0%   100.0%      15.293s       1.53e+01s     Py       1        1   forall_inplace,cpu,window_adaptation}
   0.0%   100.0%       0.000s       2.09e-04s     Py       1        1   SolveTriangular{lower=True, trans=0, unit_diagonal=False, check_finite=True}
   0.0%   100.0%       0.000s       5.72e-05s     Py       1        1   SolveTriangular{lower=False, trans=0, unit_diagonal=False, check_finite=True}
   0.0%   100.0%       0.000s       2.09e-06s     C       12       12   IncSubtensor{InplaceSet;:int64:}
   0.0%   100.0%       0.000s       1.00e-05s     C        2        2   InplaceDimShuffle{x,0}
   0.0%   100.0%       0.000s       1.38e-06s     C       10       10   AllocEmpty{dtype='float64'}
   0.0%   100.0%       0.000s       2.62e-06s     C        5        5   IncSubtensor{Set;:int64:}
   0.0%   100.0%       0.000s       7.15e-06s     C        1        1   InplaceDimShuffle{x}
   0.0%   100.0%       0.000s       6.20e-07s     C        5        5   Subtensor{uint8}
   0.0%   100.0%       0.000s       1.07e-06s     C        2        2   AllocEmpty{dtype='int64'}
   0.0%   100.0%       0.000s       2.15e-06s     C        1        1   Elemwise{sub,no_inplace}
   0.0%   100.0%       0.000s       1.91e-06s     C        1        1   Sum{acc_dtype=float64}
   0.0%   100.0%       0.000s       3.18e-07s     C        3        3   Unbroadcast{0}
   0.0%   100.0%       0.000s       9.54e-07s     C        1        1   Elemwise{Composite{(i0 - (i1 * i2))}}[(0, 2)]
   0.0%   100.0%       0.000s       0.00e+00s     C        1        1   Elemwise{Sqr}[(0, 0)]
   0.0%   100.0%       0.000s       0.00e+00s     C        1        1   Elemwise{Neg}[(0, 0)]
   ... (remaining 0 Ops account for   0.00%(0.00s) of the runtime)

Apply
------
<% time> <sum %> <apply time> <time per call> <#call> <id> <Apply name>
  100.0%   100.0%      15.293s       1.53e+01s      1    42   forall_inplace,cpu,window_adaptation}(TensorConstant{1000}, TensorConstant{[False Fal..lse  True]
}, TensorConstant{[  0   1  ..7 998 999]}, IncSubtensor{InplaceSet;:int64:}.0, IncSubtensor{InplaceSet;:int64:}.0, IncSubtensor{InplaceSet;:int64:}.0, IncSubte
nsor{Set;:int64:}.0, IncSubtensor{Set;:int64:}.0, IncSubtensor{Set;:int64:}.0, IncSubtensor{Set;:int64:}.0, IncSubtensor{InplaceSet;:int64:}.0, IncSubtensor{In
placeSet;:int64:}.0, IncSubtensor{Inpl
   0.0%   100.0%       0.000s       2.09e-04s      1    25   SolveTriangular{lower=True, trans=0, unit_diagonal=False, check_finite=True}(TensorConstant{[[1. 0
.]
 [0. 2.]]}, Elemwise{sub,no_inplace}.0)
   0.0%   100.0%       0.000s       5.72e-05s      1    30   SolveTriangular{lower=False, trans=0, unit_diagonal=False, check_finite=True}(TensorConstant{[[1. 
0.]
 [0. 2.]]}, SolveTriangular{lower=True, trans=0, unit_diagonal=False, check_finite=True}.0)
   0.0%   100.0%       0.000s       1.50e-05s      1    10   InplaceDimShuffle{x,0}(<TensorType(float64, (2,))>)
   0.0%   100.0%       0.000s       7.15e-06s      1    36   InplaceDimShuffle{x}(Sum{acc_dtype=float64}.0)
   0.0%   100.0%       0.000s       7.15e-06s      1    31   IncSubtensor{InplaceSet;:int64:}(AllocEmpty{dtype='float64'}.0, Unbroadcast{0}.0, ScalarConstant{1
})
   0.0%   100.0%       0.000s       5.01e-06s      1    33   InplaceDimShuffle{x,0}(SolveTriangular{lower=False, trans=0, unit_diagonal=False, check_finite=Tru
e}.0)
   0.0%   100.0%       0.000s       4.05e-06s      1    19   IncSubtensor{Set;:int64:}(AllocEmpty{dtype='int64'}.0, TensorConstant{(1,) of 1}, ScalarConstant{1
})
   0.0%   100.0%       0.000s       3.10e-06s      1    37   IncSubtensor{InplaceSet;:int64:}(AllocEmpty{dtype='float64'}.0, Unbroadcast{0}.0, ScalarConstant{1
})
   0.0%   100.0%       0.000s       3.10e-06s      1    22   IncSubtensor{Set;:int64:}(AllocEmpty{dtype='float64'}.0, TensorConstant{(1,) of 0.0}, ScalarConsta
nt{1})
   0.0%   100.0%       0.000s       3.10e-06s      1    13   AllocEmpty{dtype='float64'}(TensorConstant{2}, TensorConstant{2})
   0.0%   100.0%       0.000s       2.86e-06s      1    41   IncSubtensor{InplaceSet;:int64:}(AllocEmpty{dtype='float64'}.0, Elemwise{Neg}[(0, 0)].0, ScalarCon
stant{1})
   0.0%   100.0%       0.000s       2.15e-06s      1    23   IncSubtensor{InplaceSet;:int64:}(AllocEmpty{dtype='float64'}.0, TensorConstant{(1, 2) of 0.0}, Sca
larConstant{1})
   0.0%   100.0%       0.000s       2.15e-06s      1    14   IncSubtensor{Set;:int64:}(AllocEmpty{dtype='float64'}.0, TensorConstant{(1,) of -inf}, ScalarConst
ant{1})
   0.0%   100.0%       0.000s       2.15e-06s      1    11   AllocEmpty{dtype='float64'}(TensorConstant{2}, TensorConstant{2})
   0.0%   100.0%       0.000s       2.15e-06s      1     9   Elemwise{sub,no_inplace}(<TensorType(float64, (2,))>, TensorConstant{[0. 3.]})
   0.0%   100.0%       0.000s       1.91e-06s      1    47   Subtensor{uint8}(forall_inplace,cpu,window_adaptation}.0, ScalarConstant{1})
   0.0%   100.0%       0.000s       1.91e-06s      1    34   Sum{acc_dtype=float64}(Elemwise{Sqr}[(0, 0)].0)
   0.0%   100.0%       0.000s       1.91e-06s      1    24   IncSubtensor{InplaceSet;:int64:}(AllocEmpty{dtype='float64'}.0, TensorConstant{(1, 2) of 0.0}, Sca
larConstant{1})
   0.0%   100.0%       0.000s       1.91e-06s      1    21   IncSubtensor{Set;:int64:}(AllocEmpty{dtype='float64'}.0, TensorConstant{(1,) of 0.0}, ScalarConsta
nt{1})
   ... (remaining 28 Apply instances account for 0.00%(0.00s) of the runtime)


Scan overhead:
<Scan op time(s)> <sub scan fct time(s)> <sub scan op time(s)> <sub scan fct time(% scan op time)> <sub scan op time(% scan op time)> <node>
  One scan node do not have its inner profile enabled. If you enable Aesara profiler with 'aesara.function(..., profile=True)', you must manually enable the pr
ofiling for each scan too: 'aesara.scan(...,profile=True)'. Or use Aesara flag 'profile=True'.
  No scan have its inner profile enabled.
Here are tips to potentially make your code run faster
                 (if you think of new ones, suggest them on the mailing list).
                 Test them first, as they are not guaranteed to always provide a speedup.
  - Try the Aesara flag floatX=float32

real    3m30.446s
user    3m12.514s
sys     0m13.513s

Please provide any additional information below.

Versions and main components

• Aesara version: 2.8.12
• Aesara config (python -c "import aesara; print(aesara.config)"):

Details

floatX ({'float32', 'float64', 'float16'}) 
    Doc:  Default floating-point precision for python casts.

Note: float16 support is experimental, use at your own risk.
    Value:  float64

warn_float64 ({'pdb', 'ignore', '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 0x7fdca2cc2490>>) 
    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 ({'default', 'more'}) 
    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 cas
es.
    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 0x7fdc93c9f150>>) 
    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 0x7fdc98f5f190>>) 
    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 0x7fdc93d150d0>>) 
    Doc:  Print some global statistics (time spent) at the end
    Value:  False

assert_no_cpu_op ({'pdb', 'ignore', '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 0x7fdc93c9f450>>) 
    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 0x7fdc93c9f510>
    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++ cod
e is compiled.
    Value:  /home/zoj/micromamba/envs/aehmc-dev/bin/g++

linker ({'vm', 'cvm_nogc', 'c|py_nogc', 'c', 'py', 'c|py', 'cvm', 'vm_nogc'}) 
    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 0x7fdc93c9f690>>) 
    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 ({'None', 'o1', 'fast_compile', 'o2', 'o3', 'unsafe', 'merge', 'fast_run', 'o4'}) 
    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 0x7fdca174c050>>) 
    Doc:  If True, we print all optimization being applied
    Value:  False

on_opt_error ({'pdb', 'ignore', '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 0x7fdc93c9f950>>) 
    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:  

cmodule__warn_no_version (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc93c9fad0>>) 
    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 asso
ciated to at least one of those Ops.
    Value:  False

cmodule__remove_gxx_opt (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc98f5f250>>) 
    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 0x7fdc93c9fcd0>>) 
    Doc:  If True, will print compilation warnings.
    Value:  False

cmodule__preload_cache (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc93c9fdd0>>) 
    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 0x7fdc93c9fd50>>) 
    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 0x7fdc93ca8250>>) 
    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 0x7fdc93ca8350>>) 
    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 0x7fdc93ca8750>>) 
    Doc:  DEPRECATED: If True, enable the experimental optimization local_alloc_elemwise. Generates error if not True. Use optimizer_excluding=local_alloc_elem
wise to disable.
    Value:  True

experimental__local_alloc_elemwise_assert (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdca2cc0ed0>>) 
    Doc:  When the local_alloc_elemwise is applied, add an assert to highlight shape errors.
    Value:  True

warn__ignore_bug_before ({'0.3', '0.5', '1.0', '0.9', '0.7', '0.8', '1.0.5', '1.0.4', '1.0.2', '1.0.3', '0.8.2', '0.4.1', 'None', '0.8.1', '0.4', 'all', '0.10'
, '1.0.1', '0.6'}) 
    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 th
e 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 ({'low', 'high'}) 
    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 wi
ll 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 0x7fdca2cc0dd0>>) 
    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 de
bugging `my_var` will print the same as `my_var.tag.test_value` when a test value is defined.
    Value:  False

compute_test_value ({'ignore', 'warn', 'raise', 'off', 'pdb'}) 
    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 hel
ps the user track down problems in the graph before it gets optimized.
    Value:  off

compute_test_value_opt ({'ignore', 'warn', 'raise', 'off', 'pdb'}) 
    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 0x7fdca0e97310>>) 
    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 red
uce the number of generated C files.
    Value:  True

NanGuardMode__nan_is_error (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdca2cc0e90>>) 
    Doc:  Default value for nan_is_error
    Value:  True

NanGuardMode__inf_is_error (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc93ca8dd0>>) 
    Doc:  Default value for inf_is_error
    Value:  True

NanGuardMode__big_is_error (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc93ca8fd0>>) 
    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 0x7fdc93ca9150>>) 
    Doc:  Run C implementations where possible
    Value:  True

DebugMode__check_py (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc93ca9290>>) 
    Doc:  Run Python implementations where possible
    Value:  True

DebugMode__check_finite (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc93ca9390>>) 
    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 0x7fdc93ca93d0>>) 
    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 0x7fdc93ca96d0>>) 
    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 0x7fdc93ca9990>>) 
    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 0x7fdc93ca9c90>>) 
    Doc:  Do a debugprint of the profiled functions
    Value:  False

profiling__ignore_first_call (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc93ca9bd0>>) 
    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 0x7fdc93ca9a50>>) 
    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 som
e 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 num
ber 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 pea
kmemory usage, triyng both could give you a small gain.
    Value:  regular

check_stack_trace ({'off', 'warn', 'raise', 'log'}) 
    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 st
ack 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 0x7fdc93caa4d0>>) 
    Doc:  If VM should collect profile information
    Value:  False

profile_optimizer (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc93caa550>>) 
    Doc:  If VM should collect optimizer profile information
    Value:  False

profile_memory (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc93caa5d0>>) 
    Doc:  If VM should collect memory profile information and print it
    Value:  False

<aesara.configparser.ConfigParam object at 0x7fdc93caa650>
    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 bein
g 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 0x7fdc93caa7d0>>) 
    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 0x7fdc93caa950>>) 
    Doc:  If True, use Numba's fastmath mode.
    Value:  True

numba__cache (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc93caa9d0>>) 
    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_version)s-%(python_bitwidth)s.
    Value:  compiledir_%(short_platform)s-%(processor)s-%(python_version)s-%(python_bitwidth)s

<aesara.configparser.ConfigParam object at 0x7fdca25add90>
    Doc:  platform-independent root directory for compiled modules
    Value:  /home/zoj/.aesara

<aesara.configparser.ConfigParam object at 0x7fdc93cab250>
    Doc:  platform-dependent cache directory for compiled modules
    Value:  /home/zoj/.aesara/compiledir_Linux-6.1--lts-x86_64-with-glibc2.37--3.11.3-64

blas__ldflags (<class 'str'>) 
    Doc:  lib[s] to include for [Fortran] level-3 blas implementation
    Value:  -L/home/zoj/micromamba/envs/aehmc-dev/lib -lcblas -lblas -lcblas -lblas

blas__check_openmp (<bound method BoolParam._apply of <aesara.configparser.BoolParam object at 0x7fdc93922b90>>) 
    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 0x7fdc8ffbae50>>) 
    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 0x7fdc907f4c50>>) 
    Doc:  Allow/disallow memory preallocation for outputs inside of scan (default: True)
    Value:  True

• Python version: 3.11.3
• Operating system: Linux
• How did you install Aesara: conda

[brandonwillard]

Profiling graphs with Scans (or any other Ops that hold Aesara graphs inside them) is currently complicated by the way in which their inner-graphs are compiled. Simply put, the inner-graphs of Scans are compiled when a thunk is created for an Apply node that uses said Scan (see here).

This means that all the rewriting and compilation/thunk creation for inner-graphs occurs separately and in a weird lazy, nested fashion (e.g. imagine Scans containing other Scans, as is the case here). These profiling results only show one layer of that compilation stack, but the total times reflect all the rewriting and compilation in the stack.

aesara-devs/aesara#824 attempts to "flatten" all this out so that—for instance—all the timing results can be presented coherently in one profile. Also, there are a few big rewriting issues that are caused by this undesirable nesting, and aesara-devs/aesara#824 attempts to address those. I have some updates to push there, but we may want/need to move our focus to those changes in order to address this issue more easily (or at all).

In the meantime, as the output profile output implies, we might need to hack up a means of setting profile=True in all the Scans constructed in this example. With that, we can un-nest the profile results and try to figure out exactly where all the time is spent.

The top layer profile output you provided definitively tells us that linking/thunk creation/compilation is slow. Now, we need to know if that's due to compilation using the C backend and/or rewriting. We can use config.profile_optimizer to help distinguish between the two to start.

[brandonwillard]

I've used something like the following in the past to read profiling info on Scans:

scan_fn = aesara.function(...)

for node in scan_fn.maker.fgraph.apply_nodes:
    if isinstance(node.op, Scan) and hasattr(node.op.fn, "profile"):
        node.op.fn.profile.summary()

That will only get first-layer Scans, though; one needs to descend into the inner-graphs of each Scan in order to get them all. aesara-devs/aesara#824 contains a Feature that tracks all the inner-graphs added to a FunctionGraph. We can merge that separately and use it if it helps.