Stream: helpdesk (published)

Topic: Check if a struct type is "fully generic"

Brian Chen (Feb 11 2023 at 16:38):

By which I mean it should be possible to construct a new instance with each field set to any type I want. So these would pass the test:

struct S1{A, B}
  a::A
  b::B
end

struct S2{A}
  a::A
  b::Any
end

But these would not because they put additional constraints on field types:

struct S3{A}
  a::A
  b::A
end

struct S4{A, B<:Vector{A}}
  a::A
  b::B
end

Sukera (Feb 11 2023 at 17:06):

all(==(Any), fieldtypes(foo)) can at least prevent S4

Sukera (Feb 11 2023 at 17:07):

I don't think we have an interface for querying whether two fields of a parametrized struct are constrained to be the same type

Sukera (Feb 11 2023 at 17:08):

after all, you can do S3{Any} and still put whatever you into it, you just lose type stability

Mason Protter (Feb 11 2023 at 18:08):

Yeah I don’t believe we have a way to do this.

Brian Chen (Feb 11 2023 at 18:13):

The fieldtypes trick is a good one. I didn't realize it defaulted to Any for parametric types

Sukera (Feb 11 2023 at 18:15):

what else should it default to? All field type limitations due to type parameters are just refinements of Any after all

Brian Chen (Feb 11 2023 at 18:17):

Some custom marker perhaps. Either way, less important since it doesn't cover every case

Brian Chen (Feb 11 2023 at 18:17):

Context here is that I'm trying to figure out whether a Ref{Any} containing a NamedTuple, a NamedTuple of Refs or a mutable struct which allows all fields to be independently typed is more efficient

Sukera (Feb 11 2023 at 18:18):

more efficient for doing what?

Brian Chen (Feb 11 2023 at 18:18):

Storing accumulated values over time

Sukera (Feb 11 2023 at 18:19):

my gut says that the raw NamedTuple will be the same as a mutable struct with all types parametrized (that's more or less what it is already)

Sukera (Feb 11 2023 at 18:19):

except less convenient, due to not being able to set individual fields

Sukera (Feb 11 2023 at 18:19):

but storing things continously over time really screams "struct of vectors" to me

Brian Chen (Feb 11 2023 at 18:20):

No history is required, would just be pure memory overhead. This is for accumulating gradients of arbitrary mutable structs

Brian Chen (Feb 11 2023 at 18:22):

My main knowledge gap right now is knowing how each type of value is stored. Specifically what and how much is boxed

Sukera (Feb 11 2023 at 18:25):

• If the type of the field is non-concrete -> Box (most of the time, but the refinement is kept so it doesn't fall back to Any)
• If the fieldtype is not isbitstype (or otherwise the object has a size that's not statically inferrable from the type) -> it's a pointer (also with type safety)

Sukera (Feb 11 2023 at 18:25):

making this work in a generic way is difficult, and with a struct on your end doubly so, since they can't hold a variable number of fields

Brian Chen (Feb 11 2023 at 18:26):

That's why you repurpose the existing struct. But to do that you need to know whether you can plug the types you want into its fields, hence the original question

Brian Chen (Feb 11 2023 at 18:27):

*struct type. We can assume only mutable structs are being considered here

Brenhin Keller (Feb 11 2023 at 19:05):

What's the application for this?

Brian Chen (Feb 11 2023 at 22:07):

@Brenhin Keller I was thinking about whether Zygote's current method for tracking gradients of mutable structs could be improved meaningfully. Right now the mechanism is a NamedTuple wrapped in a Ref{Any}. Any update has to deref the ref, update the NamedTuple out-of-place and write it back. Rampant type instabilities aside, this seems pretty inefficient.

Brian Chen (Feb 11 2023 at 22:11):

The problem is that most of the going alternatives seem to have tradeoffs. Using a NamedTuple of Refs like https://github.com/MasonProtter/MutableNamedTuples.jl means O(fields) allocations even as we're trying to get rid of those by removing type instabilities. Creating a copy of the original struct with tangent types instead of primal ones would be great, but that requires knowing if it's safe to do so (hence this thread).

Brian Chen (Feb 11 2023 at 22:13):

Currently I'm leaning towards defining an internal type which works like MArray / https://github.com/JuliaLang/julia/pull/35453. More complexity but should only require one allocation per instance and not incur double indirects on field reads/writes

Brenhin Keller (Feb 11 2023 at 22:26):

Oh interesting

Brenhin Keller (Feb 11 2023 at 22:27):

Yeah that sounds better

Mason Protter (Feb 11 2023 at 23:05):

Why not just do something like

mutable struct AdjointStruct{T, fields, FieldTangentTypes}
    fields::NamedTuple{fields, FieldTangentTypes}
end

?

Mason Protter (Feb 11 2023 at 23:07):

I don't get why you'd want to wrap it in a Ref{Any}

Brian Chen (Feb 11 2023 at 23:26):

That's basically what I'm proposing

Brian Chen (Feb 11 2023 at 23:27):

The Ref{Any} is how things are done now

Brian Chen (Feb 19 2023 at 22:11):

After looking into this more, it may be dead end because of https://github.com/JuliaArrays/StaticArrays.jl/blob/d6e0fde34e5b3f02a7399d89cbed9c8b6e036831/src/MArray.jl#L37-L39 :(

Brian Chen (Feb 19 2023 at 22:30):

Maybe O(fields) small allocations using a MutableNamedTuples-like struct wouldn't be too bad if they happen to be pooled with good locality.

Mason Protter (Feb 20 2023 at 00:13):

@Brian Chen You can avoid that using something like Accessors.jl

Mason Protter (Feb 20 2023 at 00:14):

using Accessors, ConstructionBase

mutable struct AdjointStruct{T, fields, FieldTangentTypes}
    fields::NamedTuple{fields, FieldTangentTypes}
end
function Base.setproperty!(s::AdjointStruct{T, fieldnames, FTs}, name::Symbol, x) where {T, fieldnames, FTs}
    fields = s.fields
    newfields = set(fields, Accessors.opticcompose((PropertyLens){:x}()), x)
    setfield!(s, :fields, newfields)
end

ConstructionBase.constructorof(::Type{T}) where {T<: AdjointStruct} = T

Mason Protter (Feb 20 2023 at 00:14):

julia> let s = AdjointStruct{Foo, (:x, :y), Tuple{String, Symbol}}((;x="hi", y=:bye))
          @btime $s.x = "bye"
          s
       end
  1.940 ns (0 allocations: 0 bytes)
AdjointStruct{Foo, (:x, :y), Tuple{String, Symbol}}((x = "bye", y = :bye))

Brian Chen (Feb 20 2023 at 00:31):

Isn't that more or less the Ref approach but with a nicer interface? At least looking at @code_llvm, I see unpacking of every field in s.fields and then re-packing, all for what should be a write to a single field.

Brian Chen (Feb 20 2023 at 00:32):

(this is the case for using a type stable Ref as well)

Last updated: Oct 02 2023 at 04:34 UTC