Deeponet multi-output fix

[ayushinav]

Closes #9

[avik-pal]

Remove these, printing was fixed upstream

[KirillZubov]

input for additional layer should be size of inner embedding size

[KirillZubov]

So it not need reduction/sum/dropdim before additional layer. It should be additional = Chain(Dense(16 => 4)); here. Otherwise It's created a bottleneck and we lose information here.

[ayushinav]

Should be fixed now. Using the linear layer as additional layer for the cases where we do not have the additional layer did not seem ideal to me because it would imply weighted sum, where the weights would be learnt during training, but since DeepONets by default take the dot product, aka non-weighted sum, which could be required by many users.

[ayushinav]

Other than the doctests, the failing test cases were because the compiler because Tuple{Array{Float32, 4}, ...} would not be a subtype of Tuple{Union{Array{Float32, 3}, Array{Float32, 4}}

return type 
Tuple{Array{Float32, 4}, @NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}, additional::@NamedTuple{}}} 
does not match inferred return type 
Tuple{Union{Array{Float32, 3}, Array{Float32, 4}}, @NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}, additional::@NamedTuple{}}}

The tests pass if I comment out the @inferred and @jet cases.

[avik-pal]

Tuple{Union{Array{Float32, 3}, Array{Float32, 4}}, @NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}, additional::@NamedTuple{}}}

This is bad. You are returning a 3D array / 4D array based on the input sizes (which won't be type inferred). Avoid doing the dropdims

[ayushinav]

You are returning a 3D array / 4D array based on the input sizes (which won't be type inferred). Avoid doing the dropdims

Not sure how this might be an issue because the scalar tests also have the same dropdims calling

@inline function __project(b::AbstractArray{T1, 2}, t::AbstractArray{T2, 3},
        additional::Nothing) where {T1, T2}
    # b : p x nb
    # t : p x N x nb
    b_ = reshape(b, size(b, 1), 1, size(b, 2)) # p x 1 x nb
    return dropdims(sum(b_ .* t; dims=1); dims=1) # N x nb
end

still pass the test, and Scalar II and Vector Additonal layer tests calling

@inline function __project(
        b::AbstractArray{T1, 3}, t::AbstractArray{T2, 3}, additional::T) where {T1, T2, T}
    # b : p x u x nb
    # t : p x N x nb

    if size(b, 2) == 1 || size(t, 2) == 1
        return additional(b .* t) # p x N x nb => out_dims x N x nb
    else
        b_ = reshape(b, size(b)[1:2]..., 1, size(b, 3)) # p x u x 1 x nb
        t_ = reshape(t, size(t, 1), 1, size(t)[2:end]...) # p x 1 x N x nb

        return additional(b_ .* t_) # p x u x N x nb => out_size x N x nb
    end
end

fail.

[ayushinav]

Only doc tests fail for now.

[avik-pal]

Rebase with the latest changes to main.

[avik-pal]

set your git config to rebase on pull instead of merge, else the commit history gets royally messed up.

[codecov]

Codecov Report

Attention: Patch coverage is 83.33333% with 5 lines in your changes missing coverage. Please review.

Project coverage is 93.20%. Comparing base (aaf7d45) to head (a8149e2).

Files	Patch %	Lines
src/utils.jl	80.76%	5 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (aaf7d45) and HEAD (a8149e2). Click for more details.

HEAD has 1 upload less than BASE

Flag	BASE (aaf7d45)	HEAD (a8149e2)
	5	4

Additional details and impacted files

@@             Coverage Diff             @@
##              main      #11      +/-   ##
===========================================
- Coverage   100.00%   93.20%   -6.80%     
===========================================
  Files            7        7              
  Lines           77      103      +26     
===========================================
+ Hits            77       96      +19     
- Misses           0        7       +7     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.