bayespecon._ops.SparseFlowSolveOp¶

class bayespecon._ops.SparseFlowSolveOp(Wd, Wo, Ww)[source]¶

Differentiable sparse solve \(\eta = A(\rho)^{-1} b\).

Wraps scipy.sparse.linalg.spsolve() as a pytensor Op with analytically exact first-order gradients derived via the adjoint method.

The system matrix is:

\[A(\rho_d, \rho_o, \rho_w) = I_N - \rho_d W_d - \rho_o W_o - \rho_w W_w\]

where \(W_d = I_n \otimes W\), \(W_o = W \otimes I_n\), \(W_w = W \otimes W\) are the Kronecker-product flow weight matrices and \(N = n^2\).

This Op is used by PoissonSARFlow to embed the implicit spatial filter on the log-mean of a Poisson observation model:

\[\begin{split}\eta &= A^{-1} X\beta \\ \lambda_{ij} &= \exp(\eta_{ij}) \\ y_{ij} &\sim \operatorname{Poisson}(\lambda_{ij})\end{split}\]

The Jacobian log-determinant \(\log|A(\rho)|\) is added separately via flow_logdet_pytensor() (identical to the Gaussian SAR flow model).

Gradient derivation¶

For a scalar loss \(L\), implicit differentiation of \(A\eta = b\) gives \(dA\, \eta + A\, d\eta = db\), so:

\[d\eta = A^{-1}(db - dA\, \eta)\]

The VJPs are:

\[\frac{\partial L}{\partial \rho_k} = g^\top \frac{\partial \eta}{\partial \rho_k} = -v^\top W_k \eta, \qquad \frac{\partial L}{\partial b} = v\]

where \(v = (A^\top)^{-1} g\) and \(g = \partial L / \partial \eta\) is the upstream gradient. See _SparseFlowVJPOp for the implementation.

Per-gradient-evaluation cost: 2 sparse direct solves (SuperLU) + 3 sparse matrix-vector products. For \(n \leq 100\) (\(N \leq 10^4\)) this is fast enough for NUTS sampling.

param Wd:: Destination weight matrix \(W_d = I_n \otimes W\).
type Wd:: scipy.sparse.csr_matrix, shape (N, N)
param Wo:: Origin weight matrix \(W_o = W \otimes I_n\).
type Wo:: scipy.sparse.csr_matrix, shape (N, N)
param Ww:: Network weight matrix \(W_w = W \otimes W\).
type Ww:: scipy.sparse.csr_matrix, shape (N, N)

Examples

>>> from bayespecon._ops import SparseFlowSolveOp
>>> from bayespecon.graph import flow_weight_matrices
>>> import pytensor.tensor as pt, pytensor
>>> wms = flow_weight_matrices(G)
>>> op = SparseFlowSolveOp(wms["destination"], wms["origin"], wms["network"])
>>> rho_d, rho_o, rho_w = pt.scalars("rho_d", "rho_o", "rho_w")
>>> b = pt.vector("b")
>>> eta = op(rho_d, rho_o, rho_w, b)
>>> fn = pytensor.function([rho_d, rho_o, rho_w, b], eta)

__init__(Wd, Wo, Ww)[source]¶

Methods

`L_op`(inputs, outputs, output_grads)	Compute the VJP via the adjoint method.
`R_op`(inputs, eval_points)	Construct a graph for the R-operator.
`__init__`(Wd, Wo, Ww)
`add_tag_trace`(thing[, user_line])	Add tag.trace to a node or variable.
`do_constant_folding`(fgraph, node)	Determine whether or not constant folding should be performed for the given node.
`grad`(inputs, output_grads)	Construct a graph for the gradient with respect to each input variable.
`infer_shape`(fgraph, node, input_shapes)
`inplace_on_inputs`(allowed_inplace_inputs)	Try to return a version of self that tries to inplace in as many as allowed_inplace_inputs.
`make_node`(rho_d, rho_o, rho_w, b)	Construct an Apply node that represent the application of this operation to the given inputs.
`make_py_thunk`(node, storage_map, ...[, debug])	Make a Python thunk.
`make_thunk`(node, storage_map, compute_map, ...)	Create a thunk.
`perform`(node, inputs, outputs)	Calculate the function on the inputs and put the variables in the output storage.
`prepare_node`(node, storage_map, compute_map, ...)	Make any special modifications that the Op needs before doing `Op.make_thunk()`.

Attributes

`default_output`	An `int` that specifies which output `Op.__call__()` should return.
`destroy_map`	A `dict` that maps output indices to the input indices upon which they operate in-place.
`itypes`
`otypes`
`view_map`	A `dict` that maps output indices to the input indices of which they are a view.

L_op(inputs, outputs, output_grads)[source]¶

Compute the VJP via the adjoint method.

Delegates to _SparseFlowVJPOp, which performs:

Forward re-solve: \(\eta = A^{-1} b\).
Adjoint solve: \(v = (A^\top)^{-1} g\).
Sensitivity scalars: \(\partial L / \partial \rho_k = -v^\top W_k \eta\).
Gradient w.r.t. \(b\): \(v\).

Parameters:¶

inputs : list of TensorVariable¶: [rho_d, rho_o, rho_w, b].
outputs : list of TensorVariable¶: [eta] (symbolic forward output; not used directly here).
output_grads : list of TensorVariable¶: [g] where \(g = \partial L / \partial \eta\).

Returns:¶

[grad_rho_d, grad_rho_o, grad_rho_w, grad_b].

Return type:¶

list of TensorVariable

R_op(inputs, eval_points)[source]¶

Construct a graph for the R-operator.

This method is primarily used by Rop.

Parameters:¶

inputs¶: The Op inputs.
eval_points¶: A Variable or list of Variables with the same length as inputs. Each element of eval_points specifies the value of the corresponding input at the point where the R-operator is to be evaluated.

Return type:¶

rval[i] should be Rop(f=f_i(inputs), wrt=inputs, eval_points=eval_points).

static add_tag_trace(thing, user_line=None)[source]¶

Add tag.trace to a node or variable.

The argument is returned after being affected (inplace).

Parameters:¶

thing¶: The object where we add .tag.trace.
user_line=None¶: The max number of user line to keep.

Notes

We also use config.traceback__limit for the maximum number of stack level we look.

default_output = None[source]¶

An int that specifies which output Op.__call__() should return. If None, then all outputs are returned.

A subclass should not change this class variable, but instead override it with a subclass variable or an instance variable.

destroy_map = {}[source]¶

A dict that maps output indices to the input indices upon which they operate in-place.

Examples

destroy_map = {0: [1]} # first output operates in-place on second input
destroy_map = {1: [0]} # second output operates in-place on first input

do_constant_folding(fgraph, node)[source]¶

Determine whether or not constant folding should be performed for the given node.

This allows each Op to determine if it wants to be constant folded when all its inputs are constant. This allows it to choose where it puts its memory/speed trade-off. Also, it could make things faster as constants can’t be used for in-place operations (see *IncSubtensor).

Parameters:¶

node : Apply¶: The node for which the constant folding determination is made.

Returns:¶

res

Return type:¶

bool

grad(inputs, output_grads)[source]¶

Construct a graph for the gradient with respect to each input variable.

Each returned Variable represents the gradient with respect to that input computed based on the symbolic gradients with respect to each output. If the output is not differentiable with respect to an input, then this method should return an instance of type NullType for that input.

Using the reverse-mode AD characterization given in [1], for a \(C = f(A, B)\) representing the function implemented by the Op and its two arguments \(A\) and \(B\), given by the Variables in inputs, the values returned by Op.grad represent the quantities \(\bar{A} \equiv \frac{\partial S_O}{A}\) and \(\bar{B}\), for some scalar output term \(S_O\) of \(C\) in

\[\operatorname{Tr}\left(\bar{C}^\top dC\right) = \operatorname{Tr}\left(\bar{A}^\top dA\right) + \operatorname{Tr}\left(\bar{B}^\top dB\right)\]

Parameters:¶

inputs¶: The input variables.
output_grads¶: The gradients of the output variables.

Returns:¶

The gradients with respect to each Variable in inputs.

Return type:¶

grads

References

infer_shape(fgraph, node, input_shapes)[source]¶

inplace_on_inputs(allowed_inplace_inputs)[source]¶: Try to return a version of self that tries to inplace in as many as allowed_inplace_inputs.

itypes = None[source]¶

make_node(rho_d, rho_o, rho_w, b)[source]¶

Construct an Apply node that represent the application of this operation to the given inputs.

This must be implemented by sub-classes.

Returns:¶: node – The constructed Apply node.
Return type:¶: Apply

make_py_thunk(node, storage_map, compute_map, no_recycling, debug=False)[source]¶

Make a Python thunk.

Like Op.make_thunk() but only makes Python thunks.

make_thunk(node, storage_map, compute_map, no_recycling, impl=None)[source]¶

Create a thunk.

This function must return a thunk, that is a zero-arguments function that encapsulates the computation to be performed by this op on the arguments of the node.

Parameters:¶

node¶: Something previously returned by Op.make_node().
storage_map¶: A dict mapping Variables to single-element lists where a computed value for each Variable may be found.
compute_map¶: A dict mapping Variables to single-element lists where a boolean value can be found. The boolean indicates whether the Variable’s storage_map container contains a valid value (i.e. True) or whether it has not been computed yet (i.e. False).
no_recycling¶: List of Variables for which it is forbidden to reuse memory allocated by a previous call.
impl : str¶: Description for the type of node created (e.g. "c", "py", etc.)

Notes

If the thunk consults the storage_map on every call, it is safe for it to ignore the no_recycling argument, because elements of the no_recycling list will have a value of None in the storage_map. If the thunk can potentially cache return values (like CLinker does), then it must not do so for variables in the no_recycling list.

Op.prepare_node() is always called. If it tries 'c' and it fails, then it tries 'py', and Op.prepare_node() will be called twice.

otypes = None[source]¶

perform(node, inputs, outputs)[source]¶

Calculate the function on the inputs and put the variables in the output storage.

Parameters:¶

node¶: The symbolic Apply node that represents this computation.
inputs¶: Immutable sequence of non-symbolic/numeric inputs. These are the values of each Variable in node.inputs.
output_storage: List of mutable single-element lists (do not change the length of these lists). Each sub-list corresponds to value of each Variable in node.outputs. The primary purpose of this method is to set the values of these sub-lists.

Notes

The output_storage list might contain data. If an element of output_storage is not None, it has to be of the right type, for instance, for a TensorVariable, it has to be a NumPy ndarray with the right number of dimensions and the correct dtype. Its shape and stride pattern can be arbitrary. It is not guaranteed that such pre-set values were produced by a previous call to this Op.perform(); they could’ve been allocated by another Op’s perform method. An Op is free to reuse output_storage as it sees fit, or to discard it and allocate new memory.

prepare_node(node, storage_map, compute_map, impl)[source]¶

Make any special modifications that the Op needs before doing Op.make_thunk().

This can modify the node inplace and should return nothing.

It can be called multiple time with different impl values.

Warning

It is the Op’s responsibility to not re-prepare the node when it isn’t good to do so.

view_map = {}[source]¶

A dict that maps output indices to the input indices of which they are a view.

Examples

view_map = {0: [1]} # first output is a view of second input
view_map = {1: [0]} # second output is a view of first input