Can NeuralPDE use neural nets with more than 1 outputs?


In the following code, I would like to use one neural net with 3 outputs instead of 3 neural nets with 1 output each. I got errors when I try. Is it possible ? All the example I have seen use different nets for different functions. But I think it would be more effective to use one net with more outputs.
What do I need to change to make it possible ?

using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers, Random
import ModelingToolkit: Interval
import Dates
using GLMakie
@parameters x,y,z


domains = [x ∈ Interval(-L, +L),
           y ∈ Interval(-L, +L),
	   z ∈ Interval(-L, +L)]

@variables Fx(..),Fy(..),Fz(..)

function G(x,y,z)

Dx = Differential(x)
Dy = Differential(y)
Dz = Differential(z)


bcs = [Fx(0,0,0)~0,Fy(0,0,0)~0,Fz(0,0,0)~0]

input_ = length(domains)
n = 12

# chain = Lux.Chain(Dense(input_, n, Lux.asinh), Dense(n, n, Lux.asinh),Dense(n, n, Lux.asinh), Dense(n, 3))

chain = [Lux.Chain(Dense(input_, n, Lux.asinh), Dense(n, n, Lux.asinh),Dense(n, n, Lux.asinh), Dense(n, 1)) for _ in 1:3]

ps, st = Lux.setup(rng, chain)

strategy = QuadratureTraining()

discretization = PhysicsInformedNN(chain, strategy)

@named pdesystem = PDESystem(eqs, bcs, domains, [x,y,z], [Fx(x,y,z),Fy(x,y,z),Fz(x,y,z)])
prob = discretize(pdesystem, discretization)
sym_prob = symbolic_discretize(pdesystem, discretization)

pde_inner_loss_functions = sym_prob.loss_functions.pde_loss_functions
bcs_inner_loss_functions = sym_prob.loss_functions.bc_loss_functions


callback = function (p, l)
    if 0==i%mod
    print("$i loss: ", l)
    print(" pde: ", map(l_ -> l_(p), pde_inner_loss_functions))
    println(" bcs: ", map(l_ -> l_(p), bcs_inner_loss_functions))
    global i+=1
    return false

res = Optimization.solve(prob, Adam(0.1); callback = callback, maxiters = inter )

minimum = [res.u.depvar[sym_prob.depvars[i]] for i in 1:length(chain)]
d =
xs, ys, zs = [range(infimum(d.domain),supremum(d.domain),length=6) for d in domains]
ns = [Point3f(x,y,z) for x in xs for y in xs for z in zs]

fp(p) = Point3f(first(discretization.phi[1](p[1:3], minimum[1])),first(discretization.phi[2](p[1:3], minimum[2])),first(discretization.phi[3](p[1:3], minimum[3])))


It’s possible, but because of the way the automatic differentiation works it’s just not faster because the AD in the PDE definition is based on the outputs, and thus there’s a lot of repeated work if you use vector outputs. For this reason we use different scalar outputs. But, there is a use for the vector outputs which is still reserved and not used yet, and that’s for vector calculus which we will get to.

But say I don’t mind the repeated work then how should I do it ??