pymc-devs / pymc-devs/pytensor
Allow lists as inputs to variables and equations in `pytensor.optimize.root`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
Currently the following program is not allowed:
a, b, c, d, e, f, x, y = pt.dscalars('a b c d e f x y'.split())
eq_1 = a * x ** 2 - b * y - c
eq_2 = d * x - e * y ** 2 + f
solution, success = root(equations=[eq_1, eq_2], variables=[x, y],
method='hybr',
optimizer_kwargs={'tol':1e-8})
This has an easy work-around:
a, b, c, d, e, f = pt.dscalars('a b c d e f'.split())
xy = pt.dvector('xy', shape=(2,))
equations = pt.stack([a * xy[0] ** 2 - b * xy[1] - c, d * xy[0] - e * xy[1] ** 2 + f])
solution, success = root(equations=equations, variables=xy,
method='hybr',
optimizer_kwargs={'tol':1e-8})
But I think it's a bit of a bummer the first one doesn't work. First, the function argument is plural, so it seems to suggest that it could take a list. Second, we're making users do more work that we could probably figure out for them.
If we allow list inputs, I have 2 ideas:
- We could add special logic to handle inner functions with multiple outputs, and do a bunch of concatenation to get a single root vector and a single jacobian matrix (the latter would be similar to how we handle the args)
- We could do some graph rewriting to get everything into a canonical form, like what we do with the scalar case? Make a single
input_variable_vector, thengraph_replaceequationsto take that thing as input?
Case (2) seems better, but curious if it's too complex.
As a side note, this is something I've run into often when working with systems of equations, so some general helpers might be called for.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at the pytensor.optimize.root entry point and compare the existing scalar case with its handling of args. Inspect the proposed graph_replace approach and input_variable_vector idea, then determine how list equations and variables should be normalized. Done means the list-based example works without manual stacking and the behavior is covered by regression tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100