Chaining Custom Layers using Custom Signature

Specific Domains Machine Learning
1

Hi. Is there an idiomatic way of chaining layers in Flux when not all of the outputs of the previous layer pipe into the next layer. Unlike Flux.Chain which directly pipes all of the data, how can one create a custom layer that emulates something like:

# Pass auxilliary data/matrix A
function (model::CustomLayer)(x_0, A) 
  x_1 = model.submodel(x_0, A)
  x_2 = model.submodel(x_1, A)
  # ... 
  x_n = model.submodel(x_n_1,A)
end

Specifically, I’m looking for advice on two design choices. 1) How should I go about implementing this custom chain when all layers of the chain require the same auxilliary data. Are there Flux constructs that can express this? 2) If not and the solution just requires writing some type of loop (over n blocks), how should I properly go about this to avoid performance/type-stability issues?

Thanks.

2

Maybe the Join, Split and Parallel layers can serve your purpose. Otherwise, just define a custom layer and its forward pass;

struct CustomLayer
  ...
end

Flux.@functor CustomLayer

function (layer::CustomLayer)(x_0, A)
  ...
  return x_n
end

Using for loops is fine.

1 Like