QuantEcon / QuantEcon.py / 29061656443

94%

push

github

PERF: Compute dense Q @ v with a single gemv in bellman_operator (#857)

* PERF: Compute dense Q @ v with a single gemv in bellman_operator

Cache a 2-dimensional (n*m, n) view of Q at construction (when Q is
C-contiguous) and use it in bellman_operator, so that the expected
values are computed with one gemv call instead of a batched matmul
over n stacked (m, n) matrices; measured 2.7x on the matvec and 1.5x
on end-to-end value iteration at n=500, m=100. Non-C-contiguous Q
falls back to the previous path (regression test added). Also cache
the arange used for indexing by state in RQ_sigma and s_wise_max.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* FIX: Take the 2-d view of Q per call instead of caching it

Caching the view at construction left bellman_operator computing with
a stale array if a user rebound ddp.Q after construction (a pattern
that worked before the gemv change), while RQ_sigma and
evaluate_policy used the new one. Creating the reshape view per call
costs nothing measurable next to the gemv and removes the staleness
class entirely. Strengthen the non-contiguous-Q test with a nonuniform
Q and a direct bellman_operator/greedy-policy comparison, and add a
regression test for Q rebinding.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

42 of 42 new or added lines in 2 files covered. (100.0%)

8194 of 8751 relevant lines covered (93.64%)

0.94 hits per line