[ANN] FlexibleFunctors.jl: Functors with Per-Instance Parameters

Package Announcements
,
1

FlexibleFunctors.jl provides a utility for destructuring, restructuring, and mapping operations over a flexible set of parameters for any type.

Parameter tagging is flexible in its implementation, permitting either global or unique-per-instance parameters. Parameter extraction is recursive, extracting tagged parameters from within nested tagged fields to a flat vector representation along with a function to reconstruct the original struct. These flat representations can be modified and then reconstructed to alter the underlying parameter values.

Our syntax and inspiration draw from Functors.jl. Our implementation differs in that parameters mapped to a flat vector or NamedTuple representation can vary betweens instances of a struct.

Usage

The core usage of FlexibleFunctors.jl lies in destructure. Given an object, destructure returns a flat vector representation of the tagged parameters from that instance and a function to reconstruct that object. This reconstruction function replaces all parameters with those from its input and leaves other fields untouched.

FlexibleFunctors.jl’s flexibility comes from a small function to extend with methods for your types, parameters(m::MyType). The role of this function is to return parameter fields as a Tuple of Symbols from the given instance m. These symbols can be constant (reproducing behavior akin to Functors.jl), or they can rely on the information stored within m itself. A params fields within a struct can be used to store parameters on and instance-by-instance case, or parameters can inspect the types of m’s fields to determine parameters, given rise to ModelParameters.jl-like usage.

Functions can be mapped over all of the parameters using the fieldmap export. The syntax is similar to traditional map: fieldmap(f, x) maps f over every parameter within x and updates these values in the returned struct.

In tandem with SetField.jl, we can update a given struct’s parameters in-place. If parameters are annotated within a field of a struct, SetField.jl can update these parameters without needing to manually reconstruct a complicated model.

Example

As an example, we consider a problem of summing the areas of squares in a collection. Some squares have their width as a parameter, but one doesn’t. These Squares are compiled into Models, which themselves can tag their constituent elements as containing parameters. We are able to restructure the Squares nested within the Model, and use these resulting structs as normal.

using FlexibleFunctors, Zygote
import FlexibleFunctors: parameters

struct Model
    elements
    params
end
parameters(m::Model) = m.params

struct Square
    w
    params
end
parameters(s::Square) = s.params
area(s::Square) = s.w^2

# Squares 1 and 3 tag w as a parameter, 2 has no parameters
s1 = Square(1.0, (:w,))
s2 = Square(2.0, ())
s3 = Square(3.0, (:w,))
# Model 1 tags elements as containing parameters, Model 2 has none
m1 = Model([s1, s2, s3], (:elements,))
m2 = Model([s1, s2, s3], ())

p1, re1 = destructure(m1)
p2, re2 = destructure(m2)

julia> p1
# [1.0, 3.0]

julia> p2
# Any[]

# Replace first and third widths with new values
julia> mapreduce((s)->area(s), +, re1([3.0, 4.0]).elements)
# 29

# Uses original widths since `m2` has no parameters indicated
julia> mapreduce((s)->area(s), +, re2([3.0, 4.0]).elements)
# 14

# Map vector of parameters to an object for optimization
grad = Zygote.gradient((x) -> mapreduce(area, +, x.elements), re1([3.0, 4.0]));

julia> getfield.(grad[1].elements, :w)
# [6.0, 4.0, 8.0]

Caveats

For FlexibleFunctors.jl, reconstruction is not guaranteed to be type-stable. If the values in a reconstruction vector differ in type from the original values, no error will necessarily be thrown if they are compatible with the underlying type. Care should be taken with respect to replacing Floats with Ints and similar cases. Comparable packages below treat this case more deliberately than we have.

We currently don’t perform any checks on the argument to a reconstruction function. If your reconstruciton vector is longer than the number of parameters being reconstructed, the first elements of the vector will be used for their respective parameters and later elements will be ignored.

Related Work

  • Functors.jl served as primary inspiration for this package. We wanted the ability to easily instantiate complex structures from a flat vector representation, but we needed individual instances of any given struct to have different parameter fields. As it stands, Functors.jl is global across all instances of any given struct.
  • ParameterHandling.jl focuses on constraining the parameters and manually extending flatten methods. FlexibleFunctors.jl will automatically extract parameters once they’ve been annotated within a struct, and parameters methods should easily customize to your suit your needs.
  • ModelingParameters.jl is a great package which implements a very nice tool where Param types denote parameters within a struct. These Params are used to tag parameters, as opposed to our method where an external method operates on a struct to determine parameters with customizable behavior.
7 Likes
2

Congrats on the release!

I just wanted to clarify that this kind of dynamic parameter flexibility is possible in Functors.jl, but requires some more low-level plumbing:

struct Square
    w
    params
end

function Functors.functor(s::Square)
   function re(y)
        all_args = map(fieldnames(T)) do fn
            field = fn in s.params ? getfield(y, fn) : getfield(x, fn)
            return field
        end
        return constructorof(T)(all_args...)
    end
    func = (; (p => getfield(x, p) for p in s.params)...)
    return func, re
end

Essentially, substituting in the guts of ffunctor.

Given how much overlap there is between both codebases, I’d propose possibly integrating this functionality into Functors itself, either as a convenience method or macro (e.g. @flexiblefunctor). This would reduce the amount of fragmentation in parameter tracking packages and allow flexible functors to work with packages like Flux out of the box. WDYT?

9 Likes
3

Thanks for your feedback!

I think that integration into Functors.jl is a good solution. Given the similarity between the two, wrapping this in as something like a @flexiblefunctor macro would be ideal for proliferating the capability and de-duplicating work. The majority of use we’ve had with this has been storing parameters within an internal field, so something like the

@flexiblefunctor Square params

syntax would be great.

I’ll work on putting together a PR for Functors.jl to see if it’s an appropriate fit

6 Likes
6

Self advertising - ValueShapes would fit into that list as well. :slight_smile: (It’s currently still limited to scalars, arrays and named tuples, though).

1 Like