Benchmarking and Optimizing Aesara Code #1174
Replies: 4 comments 16 replies
-
|
Out of curiosity does the Just do Regarding PyMC, keep in mind that the gradient probably dominates the sampling time, so it can be easily be worse than what you would expect from examining the logp performance alone. I read somewhere that the Scan grad is 4-6x slower than the value evaluation, but that obviously depends on the graph itself. |
Beta Was this translation helpful? Give feedback.
-
|
Keeping everything as above and using mode="NUMBA" raises an error about a duplicate argument name in function definition. Maybe related to the None's I pass into the SyntaxError: duplicate argument 'auto_446096' in function definition
Traceback (most recent call last):
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/IPython/core/interactiveshell.py:3398 in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
Input In [67] in <cell line: 9>
matrix_kf = aesara.function([a_data, a_x0, a_P0, a_T, a_Z, a_R, a_H, a_Q], filter_result, mode='NUMBA')
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/compile/function/__init__.py:317 in function
fn = pfunc(
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/compile/function/pfunc.py:374 in pfunc
return orig_function(
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/compile/function/types.py:1763 in orig_function
fn = m.create(defaults)
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/compile/function/types.py:1656 in create
_fn, _i, _o = self.linker.make_thunk(
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/link/basic.py:254 in make_thunk
return self.make_all(
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/link/basic.py:698 in make_all
thunks, nodes, jit_fn = self.create_jitable_thunk(
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/link/basic.py:642 in create_jitable_thunk
converted_fgraph = self.fgraph_convert(
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/link/numba/linker.py:10 in fgraph_convert
return numba_funcify(fgraph, **kwargs)
File ~/opt/anaconda3/envs/econ/lib/python3.10/functools.py:889 in wrapper
return dispatch(args[0].__class__)(*args, **kw)
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/link/numba/dispatch/basic.py:381 in numba_funcify_FunctionGraph
return fgraph_to_python(
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/link/utils.py:741 in fgraph_to_python
compiled_func = op_conversion_fn(
File ~/opt/anaconda3/envs/econ/lib/python3.10/functools.py:889 in wrapper
return dispatch(args[0].__class__)(*args, **kw)
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/link/numba/dispatch/scan.py:153 in numba_funcify_Scan
scalar_op_fn = compile_function_src(
File ~/opt/anaconda3/envs/econ/lib/python3.10/site-packages/aesara/link/utils.py:605 in compile_function_src
mod_code = compile(src, filename, mode="exec")
File /var/folders/wy/ph4j9vrx23v000gc9flt78y40000gn/T/tmpmmqtt8nb:2
def scan(n_steps, auto_448050, auto_452429, auto_452428, auto_446096, auto_446096, auto_446096, auto_445347, auto_445348, auto_445350, auto_446406, auto_446412, auto_446414):
^
SyntaxError: duplicate argument 'auto_446096' in function definitionI thought about the gradient problem as well, and I profile and examine the functions created by |
Beta Was this translation helpful? Give feedback.
-
|
We should create an issue for that; it looks distinct from the other known |
Beta Was this translation helpful? Give feedback.
-
|
Playing a bit with this it seems that multiple "map outputs" (None values) is the cause of the error, I opened an issue #1176 |
Beta Was this translation helpful? Give feedback.
-
|
To reply more directly to your question @ricardoV94, the graph does work, when I explicitly carry along all computations in the scan. That is, I provide no more than one The PyMC model with Numba scan
with pm.Model() as matrix_model:
x0 = pm.Normal('x0', sigma=1, size=(1,))
P0 = pm.Gamma('P0', alpha=2, beta=8, size=(1,1))
H = pm.Gamma('H', alpha=2, beta=8, size=(1,1))
Q = pm.Gamma('Q', alpha=2, beta=8, size=(1,1))
filter_result, updates = aesara.scan(matrix_kalman_step,
sequences=[data],
outputs_info=[at.zeros_like(x0), x0, at.zeros_like(P0), P0, None],
non_sequences=[T, Z, R, H, Q],
mode='NUMBA')
a_filtered, a_predicted, P_filtered, P_predicted, ll = filter_result
loglike = pm.Potential('log_likelihood', ll.sum())
idata_matrix = pm.sample()TypeError: Inconsistency in the inner graph of scan 'scan_fn' : an input and an output are associated with the same recurrent state and should have compatible types but have type 'TensorType(float64, (1, 1))' and 'TensorType(float64, (None, None))' respectively.
TypeError Traceback (most recent call last)
Input In [21], in <cell line: 1>()
13 a_filtered, a_predicted, P_filtered, P_predicted, ll = filter_result
15 loglike = pm.Potential('log_likelihood', ll.sum())
---> 17 idata_matrix = pm.sample()
File ~\miniconda3\envs\econ\lib\site-packages\pymc\sampling.py:531, in sample(draws, step, init, n_init, initvals, trace, chain_idx, chains, cores, tune, progressbar, model, random_seed, discard_tuned_samples, compute_convergence_checks, callback, jitter_max_retries, return_inferencedata, idata_kwargs, mp_ctx, **kwargs)
528 auto_nuts_init = False
530 initial_points = None
--> 531 step = assign_step_methods(model, step, methods=pm.STEP_METHODS, step_kwargs=kwargs)
533 if isinstance(step, list):
534 step = CompoundStep(step)
File ~\miniconda3\envs\econ\lib\site-packages\pymc\sampling.py:229, in assign_step_methods(model, step, methods, step_kwargs)
221 selected = max(
222 methods,
223 key=lambda method, var=rv_var, has_gradient=has_gradient: method._competence(
224 var, has_gradient
225 ),
226 )
227 selected_steps[selected].append(var)
--> 229 return instantiate_steppers(model, steps, selected_steps, step_kwargs)
File ~\miniconda3\envs\econ\lib\site-packages\pymc\sampling.py:147, in instantiate_steppers(model, steps, selected_steps, step_kwargs)
145 args = step_kwargs.get(step_class.name, {})
146 used_keys.add(step_class.name)
--> 147 step = step_class(vars=vars, model=model, **args)
148 steps.append(step)
150 unused_args = set(step_kwargs).difference(used_keys)
File ~\miniconda3\envs\econ\lib\site-packages\pymc\step_methods\arraystep.py:89, in BlockedStep.__new__(cls, *args, **kwargs)
86 step = super().__new__(cls)
87 # If we don't return the instance we have to manually
88 # call __init__
---> 89 step.__init__([var], *args, **kwargs)
90 # Hack for creating the class correctly when unpickling.
91 step.__newargs = ([var],) + args, kwargs
File ~\miniconda3\envs\econ\lib\site-packages\pymc\step_methods\metropolis.py:229, in Metropolis.__init__(self, vars, S, proposal_dist, scaling, tune, tune_interval, model, mode, **kwargs)
226 self.mode = mode
228 shared = pm.make_shared_replacements(initial_values, vars, model)
--> 229 self.delta_logp = delta_logp(initial_values, model.logp(), vars, shared)
230 super().__init__(vars, shared)
File ~\miniconda3\envs\econ\lib\site-packages\pymc\step_methods\metropolis.py:1056, in delta_logp(point, logp, vars, shared)
1052 inarray1 = tensor_type("inarray1")
1054 logp1 = pm.CallableTensor(logp0)(inarray1)
-> 1056 f = compile_pymc([inarray1, inarray0], logp1 - logp0)
1057 f.trust_input = True
1058 return f
File ~\miniconda3\envs\econ\lib\site-packages\pymc\aesaraf.py:1034, in compile_pymc(inputs, outputs, random_seed, mode, **kwargs)
1032 opt_qry = mode.provided_optimizer.including("random_make_inplace", check_parameter_opt)
1033 mode = Mode(linker=mode.linker, optimizer=opt_qry)
-> 1034 aesara_function = aesara.function(
1035 inputs,
1036 outputs,
1037 updates={**rng_updates, **kwargs.pop("updates", {})},
1038 mode=mode,
1039 **kwargs,
1040 )
1041 return aesara_function
File ~\miniconda3\envs\econ\lib\site-packages\aesara\compile\function\__init__.py:317, in function(inputs, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input)
311 fn = orig_function(
312 inputs, outputs, mode=mode, accept_inplace=accept_inplace, name=name
313 )
314 else:
315 # note: pfunc will also call orig_function -- orig_function is
316 # a choke point that all compilation must pass through
--> 317 fn = pfunc(
318 params=inputs,
319 outputs=outputs,
320 mode=mode,
321 updates=updates,
322 givens=givens,
323 no_default_updates=no_default_updates,
324 accept_inplace=accept_inplace,
325 name=name,
326 rebuild_strict=rebuild_strict,
327 allow_input_downcast=allow_input_downcast,
328 on_unused_input=on_unused_input,
329 profile=profile,
330 output_keys=output_keys,
331 )
332 return fn
File ~\miniconda3\envs\econ\lib\site-packages\aesara\compile\function\pfunc.py:362, in pfunc(params, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input, output_keys, fgraph)
359 elif isinstance(profile, str):
360 profile = ProfileStats(message=profile)
--> 362 inputs, cloned_outputs = construct_pfunc_ins_and_outs(
363 params,
364 outputs,
365 mode,
366 updates,
367 givens,
368 no_default_updates,
369 rebuild_strict,
370 allow_input_downcast,
371 fgraph=fgraph,
372 )
374 return orig_function(
375 inputs,
376 cloned_outputs,
(...)
383 fgraph=fgraph,
384 )
File ~\miniconda3\envs\econ\lib\site-packages\aesara\compile\function\pfunc.py:492, in construct_pfunc_ins_and_outs(params, outputs, mode, updates, givens, no_default_updates, rebuild_strict, allow_input_downcast, fgraph)
489 out_list = [outputs]
490 extended_outputs = out_list + additional_outputs
--> 492 output_vars = rebuild_collect_shared(
493 extended_outputs,
494 in_variables,
495 replace=givens,
496 updates=updates,
497 rebuild_strict=rebuild_strict,
498 copy_inputs_over=True,
499 no_default_updates=no_default_updates,
500 clone_inner_graphs=True,
501 )
502 input_variables, cloned_extended_outputs, other_stuff = output_vars
503 clone_d, update_d, update_expr, shared_inputs = other_stuff
File ~\miniconda3\envs\econ\lib\site-packages\aesara\compile\function\pfunc.py:222, in rebuild_collect_shared(outputs, inputs, replace, updates, rebuild_strict, copy_inputs_over, no_default_updates, clone_inner_graphs)
220 for v in outputs:
221 if isinstance(v, Variable):
--> 222 cloned_v = clone_v_get_shared_updates(v, copy_inputs_over)
223 cloned_outputs.append(cloned_v)
224 elif isinstance(v, Out):
File ~\miniconda3\envs\econ\lib\site-packages\aesara\compile\function\pfunc.py:95, in rebuild_collect_shared.<locals>.clone_v_get_shared_updates(v, copy_inputs_over)
93 if owner not in clone_d:
94 for i in owner.inputs:
---> 95 clone_v_get_shared_updates(i, copy_inputs_over)
96 clone_node_and_cache(
97 owner,
98 clone_d,
99 strict=rebuild_strict,
100 clone_inner_graphs=clone_inner_graphs,
101 )
102 return clone_d.setdefault(v, v)
File ~\miniconda3\envs\econ\lib\site-packages\aesara\compile\function\pfunc.py:95, in rebuild_collect_shared.<locals>.clone_v_get_shared_updates(v, copy_inputs_over)
93 if owner not in clone_d:
94 for i in owner.inputs:
---> 95 clone_v_get_shared_updates(i, copy_inputs_over)
96 clone_node_and_cache(
97 owner,
98 clone_d,
99 strict=rebuild_strict,
100 clone_inner_graphs=clone_inner_graphs,
101 )
102 return clone_d.setdefault(v, v)
[... skipping similar frames: rebuild_collect_shared.<locals>.clone_v_get_shared_updates at line 95 (2 times)]
File ~\miniconda3\envs\econ\lib\site-packages\aesara\compile\function\pfunc.py:95, in rebuild_collect_shared.<locals>.clone_v_get_shared_updates(v, copy_inputs_over)
93 if owner not in clone_d:
94 for i in owner.inputs:
---> 95 clone_v_get_shared_updates(i, copy_inputs_over)
96 clone_node_and_cache(
97 owner,
98 clone_d,
99 strict=rebuild_strict,
100 clone_inner_graphs=clone_inner_graphs,
101 )
102 return clone_d.setdefault(v, v)
File ~\miniconda3\envs\econ\lib\site-packages\aesara\compile\function\pfunc.py:96, in rebuild_collect_shared.<locals>.clone_v_get_shared_updates(v, copy_inputs_over)
94 for i in owner.inputs:
95 clone_v_get_shared_updates(i, copy_inputs_over)
---> 96 clone_node_and_cache(
97 owner,
98 clone_d,
99 strict=rebuild_strict,
100 clone_inner_graphs=clone_inner_graphs,
101 )
102 return clone_d.setdefault(v, v)
103 elif isinstance(v, SharedVariable):
File ~\miniconda3\envs\econ\lib\site-packages\aesara\graph\basic.py:1038, in clone_node_and_cache(node, clone_d, clone_inner_graphs, **kwargs)
1034 new_op: Optional["Op"] = cast(Optional["Op"], clone_d.get(node.op))
1036 cloned_inputs: List[Variable] = [cast(Variable, clone_d[i]) for i in node.inputs]
-> 1038 new_node = node.clone_with_new_inputs(
1039 cloned_inputs,
1040 # Only clone inner-graph `Op`s when there isn't a cached clone (and
1041 # when `clone_inner_graphs` is enabled)
1042 clone_inner_graph=clone_inner_graphs if new_op is None else False,
1043 **kwargs,
1044 )
1046 if new_op:
1047 # If we didn't clone the inner-graph `Op` above, because
1048 # there was a cached version, set the cloned `Apply` to use
1049 # the cached clone `Op`
1050 new_node.op = new_op
File ~\miniconda3\envs\econ\lib\site-packages\aesara\graph\basic.py:291, in Apply.clone_with_new_inputs(self, inputs, strict, clone_inner_graph)
289 new_node.tag = copy(self.tag).__update__(new_node.tag)
290 else:
--> 291 new_node = self.clone(clone_inner_graph=clone_inner_graph)
292 new_node.inputs = new_inputs
293 return new_node
File ~\miniconda3\envs\econ\lib\site-packages\aesara\graph\basic.py:231, in Apply.clone(self, clone_inner_graph)
228 new_op = self.op
230 if isinstance(new_op, HasInnerGraph) and clone_inner_graph: # type: ignore
--> 231 new_op = new_op.clone() # type: ignore
233 cp = self.__class__(
234 new_op, self.inputs, [output.clone() for output in self.outputs]
235 )
236 cp.tag = copy(self.tag)
File ~\miniconda3\envs\econ\lib\site-packages\aesara\scan\op.py:1487, in Scan.clone(self)
1486 def clone(self) -> "Scan":
-> 1487 res = copy(self)
1488 res.fgraph = res.fgraph.clone()
1489 return res
File ~\miniconda3\envs\econ\lib\copy.py:102, in copy(x)
100 if isinstance(rv, str):
101 return x
--> 102 return _reconstruct(x, None, *rv)
File ~\miniconda3\envs\econ\lib\copy.py:273, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
271 state = deepcopy(state, memo)
272 if hasattr(y, '__setstate__'):
--> 273 y.__setstate__(state)
274 else:
275 if isinstance(state, tuple) and len(state) == 2:
File ~\miniconda3\envs\econ\lib\site-packages\aesara\scan\op.py:913, in Scan.__setstate__(self, d)
911 self.__dict__.update(d)
912 # Ensure that the graph associated with the inner function is valid.
--> 913 self.validate_inner_graph()
File ~\miniconda3\envs\econ\lib\site-packages\aesara\scan\op.py:634, in ScanMethodsMixin.validate_inner_graph(self)
629 type_output = self.inner_outputs[inner_oidx].type
630 if (
631 type_input.dtype != type_output.dtype
632 or type_input.broadcastable != type_output.broadcastable
633 ):
--> 634 raise TypeError(
635 "Inconsistency in the inner graph of "
636 f"scan '{self.name}' : an input and an output are "
637 "associated with the same recurrent state "
638 "and should have compatible types but have "
639 f"type '{type_input}' and '{type_output}' respectively."
640 )
TypeError: Inconsistency in the inner graph of scan 'scan_fn' : an input and an output are associated with the same recurrent state and should have compatible types but have type 'TensorType(float64, (1, 1))' and 'TensorType(float64, (None, None))' respectively. |
Beta Was this translation helpful? Give feedback.
-
Printing the graphs is often good for relative comparisons (e.g. the scalar-case graph vs. the matrix-case graph). Sometimes one can see which optimizations were performed in one case but not the other, and, if those optimizations are meaningful, they could explain any large relative differences in performance. |
Beta Was this translation helpful? Give feedback.
-
Don't forget that a |
Beta Was this translation helpful? Give feedback.
-
|
I've created #1179 and #1180 for two of the issues mentioned above. |
Beta Was this translation helpful? Give feedback.
-
What can be done to cut down on the number of DimShuffles? If a formula calls for a transpose, I just need to do it no? |
Beta Was this translation helpful? Give feedback.
-
I guess the point was just to compare how much performance degrades when one switches to linalg operators, not that @jessegrabowski was actually interested in scalar matrixes right? |
Beta Was this translation helpful? Give feedback.
-
|
Yes, the point was to demonstrate, to my self I suppose, that |
Beta Was this translation helpful? Give feedback.
-
|
I was only speaking to the obvious lack of "translation invariance" between the scalar and (scalar) matrix cases. |
Beta Was this translation helpful? Give feedback.
-
|
Based on the discussion and the profiling so far, it seems the main bottleneck is in linear algebra operators. There is another implementation of the Kalman Filter that avoids computing determinants or inverting any matrices by introducing a second scan over the states of the system. Introducing this second scan reduces all the matrix-matrix operations to at most vector-vector, with no matrix inversions at all. Combined with the observed speed of the scalar implementation, especially when sampling in PyMC, I thought this might be a fruitful way to go. Here is the so-called "Univariate Kalman Filter", which has a scan-in-a-scan. It ends up being significantly slower than the linear algebra operations, which surprised me, especially since the inner scan is only over a single state (so it should only be a single operation). Univariate Kalman Filter with a scan inside a scan to avoid `at.linalg.solve` and `at.linalg.det`
def univariate_inner_filter_step(y, Z_row, sigma_H, a, P):
v = y - Z_row.dot(a)
PZT = P.dot(Z_row.T)
F = Z_row.dot(PZT) + sigma_H
K = PZT / F
a_filtered = a + K * v
P_filtered = P - at.outer(K, K) * F
ll_inner = at.log(F) + v ** 2 / F
return a_filtered, P_filtered, ll_inner
def matrix_predict2(a_filtered, P_filtered, T, R, Q):
a_predicted = T.dot(a_filtered)
P_predicted = matrix_dot(T, P_filtered, T.T) + matrix_dot(R, Q, R.T)
# Force P_predicted to be symmetric
P_predicted = 0.5 * (P_predicted + P_predicted.T)
return a_predicted, P_predicted
def univariate_kalman_step(y, a, P, T, Z, R, H, Q):
y = y[:, None]
result, updates = aesara.scan(univariate_inner_filter_step,
sequences=[y, Z, at.diag(H)],
outputs_info=[a, P, None],
name='scan_over_states',
profile=True)
a_filtered, P_filtered, ll_inner = result
a_filtered, P_filtered = a_filtered[-1], P_filtered[-1]
a_predicted, P_predicted = matrix_predict2(a_filtered, P_filtered, T=T, R=R, Q=Q)
ll = -0.5 * ((at.neq(ll_inner, 0).sum()) * N_CONST + ll_inner.sum())
return a_filtered, a_predicted, P_filtered, P_predicted, ll
filter_result, updates = aesara.scan(univariate_kalman_step,
sequences=[a_data],
outputs_info=[None, a_x0, None, a_P0, None],
non_sequences=[a_T, a_Z, a_R, a_H, a_Q],
name='univariate_filter',
profile=True)
Compiling the inner scan by itself and timing it shows that it's a very fast function, just a few microseconds (makes sense since it has only 1 step). When I profile the whole function together, though, the inner scan ends up dominating the execution time: Outer scan profile
Inner scan profile
I don't understand the discrepancy between the execution times in the two profiles. The outer scan reports that the inner scan used 6s to execute, but the profile of the inner scan says that it executed for only 2.4s. Sampling time for this set-up in PyMC is worst of all, so I guess @ricardoV94 was right that gradients of |
Beta Was this translation helpful? Give feedback.
-
I've found that the SVD approach (e.g. here) works quite well in non-trivial multivariate problems.
Don't forget to return the inner-
Regarding the "Scan overhead:" measurements:
Yeah, gradients of Even so, the main problem with |
Beta Was this translation helpful? Give feedback.
-
|
Also, there appears to be a bug somewhere: |
Beta Was this translation helpful? Give feedback.
-
Should we open an issue for that? |
Beta Was this translation helpful? Give feedback.
-
Yes, definitely, but we need an MWE. |
Beta Was this translation helpful? Give feedback.
-
Hi everyone,
I've been working on implementing a Kalman Filter in Aesara, with the objective of computing the log-likelihood of a class of time series models. The filter is pretty simple, it just recursively predicts new data using model dynamics, computes the error between the prediction and a single time-step, and then combines the prediction and observation into a optimal fused state.
In the most general form, the filter is a bunch of matrix equations, but a special case where everything is scalar can also be implemented. Here it this special case in Aesara:
Code block for the scalar kalman filter
On 100 data-points with fixed parameters, this scalar function runs extremely fast:
511 µs ± 4.4 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each).Plugging the log-likelihood into PyMC, things also sample very fast (22 seconds). Here is the function profile, as a basis for comparison:Aesara profile output for the scalar kalman filter
Question 1: Checking the profile for
kalman_filter, 98% of the execution time is spent in the scan. That makes perfect sense, since there isn't much else here. The scan is of type Py rather than C, is that correct? Is there anything else interesting in this profile output that should catch my attention? So far, everything seems fine.What I really want, though, is to implement the filter with matrices, so that it can generalize to more complex time series models. Here is the code for a Kalman filter that uses matrices:
Matrix Kalman filter code
Implementing the same scalar model using a set of 1x1 matrices gives the following time:
11.3 ms ± 150 µs per loop (mean ± std. dev. of 7 runs, 100 loops each), or a 22x slowdown. Passing the likelihood to PyMC, sampling goes from 22 seconds (scalar model) to 4 minutes (matrix model).Once again, I look at the profile for
matrix_kalman_filter, which is dominated by the scan Op, with a bit ofDot22thrown in.Matrix Kalman Filter function profile
Question 2: How should I interpret these profile times in the presence of a single dominating apply node (the scan itself)? Can I really interpret the entirety of the speed differential to the calls to Dot22, even though this accounts for only 1.1% of the time spent in execution? (Obviously there are more differences between the two implementations to be addressed below; I am trying to ask a concrete question about this single output).
Now, these are not apples-to-apples comparisons, because aside from switching from multiplication and division to dot-products and matrix inversion, there are several additional operations introduced for numerical stability. The scalar filter has 32 apply nodes, while the matrix filter has 37, so obviously I expect the code to be somewhat slower. How much slower, and whether I am leaving any optimizations on the table in the matrix form of the filter, are open questions.
To get a sense of how much slowdown I should expect when moving from scalar math to linear algebra, I ran everything in pure python, and saw a 25x slowdown, so Aesara is actually doing a bit better here. To get a sense of what optimizations are possible, I put
njitdecorators on all the functions and re-ran the speed tests. This brought times down to 9.33 µs ± 500 ns and 733 µs ± 3.13 µs for the scalar and matrix versions, respectively. So the linear algebra slows down numba the most, by 78x, but in absolute time it's still the fastest implementation.Here a little table summarizing results:
Back to Aesara, @brandonwillard told me that all efforts start with
aesara.dprint, so I checked the graph for the log-likelihood for the matrix function, but nothing jumps out at me as being "off"; there's no obvious duplicate computation or bugs happening.dprint output for matrix kalman filter log-likelihood term
So what have I learned? Linear algebra is slower than scalar multiplication. Thanks for coming to my TED talk. In all seriousness, what can be done to speed up the matrix formulation of the filter? From looking at the profile output and the
aesara.dprintoutput, I don't see any obvious paths forward. The final goal is to be able to sample using this log-likelihood function, so speed is absolutely critical. With the simplest possible model already already taking 5 minutes, it's quite dishartening. For comparison, an ARMA(1,1) model takes over 15 minutes, and this is comprised only of 2x2 matrices. Scaling gets worse from there.All the code used to make this post is available in this gist if anyone is curious. I am really hoping that I am missing something obvious. At the very least, helpful pointers about how to read the profile summary and dprint outputs will be appreciated.
Beta Was this translation helpful? Give feedback.
All reactions