Node with multiple output keys

[mateusfvp]

Hi,
In the Google Colab tutorial Part 1: Learning to solve a constrained optimization problem, the output of node sol_map is a vector x of size 2. When defining the decision variables, the vector is sliced and then the first entry is assigned to variable x1, and the second entry is assigned to x2.

# define neural architecture for the trainable solution map
func = blocks.MLP(insize=2, outsize=2,
                bias=True,
                linear_map=slim.maps['linear'],
                nonlin=nn.ReLU,
                hsizes=[80] * 4)
# wrap neural net into symbolic representation of the solution map via the Node class: sol_map(xi) -> x
sol_map = Node(func, ['a', 'p'], ['x'], name='map')

# define decision variables
x1 = variable("x")[:, [0]]
x2 = variable("x")[:, [1]]

How to define the sol_map node such that it outputs x1 and x2 separately as two scalars so we don't need to slice a vector when defining the decision variables? I tried the following

# define neural architecture for the trainable solution map
func = blocks.MLP(insize=2, outsize=2,
                bias=True,
                linear_map=slim.maps['linear'],
                nonlin=nn.ReLU,
                hsizes=[80] * 4)
# wrap neural net into symbolic representation of the solution map via the Node class: sol_map(xi) -> x
sol_map = Node(func, ['a', 'p'], ['x1','x2'], name='map')

# define decision variables
x1 = variable("x1")
x2 = variable("x2")

but I got the following error

[drgona]

Hi @mateusfvp thanks for the inquiry.

It all depends on the definition of the callable. By default, our pre-defined MLP model supports an arbitrary number of positional arguments on inputs, internally, they are all concatenated into a single tensor. This allows handling multiple symbolic input variables via the Node class.
See the base Block class definition here:
https://github.com/pnnl/neuromancer/blob/master/src/neuromancer/modules/blocks.py#L27

As for the number of returns of the MLP class, the default is one. Therefore, the expected number of symbolic outputs is one. When you define more than one, there will be an error in the computational graph, as the rest of the variables won't be defined.

There are several workarounds to have these two variables defined separately. You can either instantiate two neural nets and then wrap them separately in two node classes:

func1 = blocks.MLP(insize=2, outsize=2,  
                bias=True,  
                linear_map=slim.maps['linear'],  
                nonlin=nn.ReLU,  
                hsizes=[80] * 4)  
func2 = blocks.MLP(insize=2, outsize=2,  
                bias=True,  
                linear_map=slim.maps['linear'],  
                nonlin=nn.ReLU,  
                hsizes=[80] * 4)  
sol_map1 = Node(func1, ['a', 'p'], ['x1'], name='map')  
sol_map2 = Node(func2, ['a', 'p'], ['x2'], name='map')  

or you can define callable that has two returns, for instance

def callable(a, p):
    x = fun1(a, p)
    y = fun2(a, p)
    return x, y

or really any custom callable with multiple outputs, for instance

def callable(a, p):
    x = a+p
    y = a*p
    return x, y