pymc-devs / pymc-devs/pytensor

Allow lists as inputs to variables and equations in `pytensor.optimize.root`

Open
#1,465 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement graph rewriting
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:

  1. 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)
  2. 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, then graph_replace equations to 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.