Adjoint(A)*A for MArray and Type Stability

Specific Domains
1

I wonder how to implement Adjoint(A)*A in type stable way for MArrays. Below is a MWE and output of @code_warntype()

struct SFoo
  SGnu::SMatrix{3,3}
end 

struct MFoo
  MGnu::MMatrix{3,3}
end 

sgnu = SMatrix{3,3}(0., 0., 0., 0., 0., 0., 0., 0., 0.)
sfoo = SFoo(sgnu)

mgnu = MMatrix{3,3}(0., 0., 0., 0., 0., 0., 0., 0., 0.)
mfoo = MFoo(mgnu)

function bar(sfoo)
    sgnu = sfoo.SGnu 
    return Adjoint(sgnu)*sgnu 
end 

function bar(mfoo)
    mgnu = mfoo.MGnu 
    return Adjoint(mgnu)*mgnu 
end

Goes well for SArray

MethodInstance for bar(::SFoo)
from bar(mfoo) in Main at In[65]:6
Arguments
#self#::Core.Const(bar)
mfoo::SFoo
Locals
mgnu::Union{}
Body::Union{}
1 ─ (mgnu = Base.getproperty(mfoo, :MGnu))
│ Core.Const(:(Main.Adjoint(mgnu)))
│ Core.Const(:(%2 * mgnu))
└── Core.Const(:(return %3))

Results in Any for MArray

MethodInstance for bar(::MFoo)
from bar(mfoo) in Main at In[65]:6
Arguments
#self#::Core.Const(bar)
mfoo::MFoo
Locals
mgnu::MMatrix{3, 3}
Body::Any
1 ─ (mgnu = Base.getproperty(mfoo, :MGnu))
│ %2 = Main.Adjoint(mgnu)::Adjoint
│ %3 = (%2 * mgnu)::Any
└── return %3

2

I’m not sure why this correctly infers for the SMatrix case, but you can fix it for the MMatrix case by fully specifying your MMatrix type like so:

struct MFoo{T}
    MGnu::MMatrix{3, 3, T, 9}
end

With this, running

mgnu = MMatrix{3,3}(0., 0., 0., 0., 0., 0., 0., 0., 0.)
mfoo = MFoo(mgnu)
bar(mfoo)

will correctly infer the return type as an MMatrix. The final type parameter (the 9) is the number of elements stored for the MMatrix, and the T is the numeric type (here, a Float64)

3

The second method overwrites the first here, so you’re studying the bar(mfoo) version in both cases. It looks type stable when the argument is an SFoo because the compiler knows it will error and never return; instances of type SFoo have no field MGnu. Note that the inferred return type is Union{}, which is a type with no instances.

To implement different methods for different types, use type annotations as follows:

function bar(sfoo::SFoo)
    sgnu = sfoo.SGnu 
    return Adjoint(sgnu)*sgnu 
end 

function bar(mfoo::MFoo)
    mgnu = mfoo.MGnu 
    return Adjoint(mgnu)*mgnu 
end

To solve the type stability issue, make the fields concretely typed as shown by @JonasWickman, though I would personally prefer a solution that requires less knowledge about the implementation of MMatrix, like the following:

struct MFoo{M<:MMatrix{3,3}}
    MGnu::M
end
1 Like
4

One more thing, you should normally avoid using the constructor Adjoint. Instead, use the function adjoint, or, equivalently, the postfix operator ', as in adjoint(mgnu) * mgnu or mgnu' * mgnu.

The main difference is that adjoint knows that it is its own inverse, so adjoint(adjoint(M)) === M. Meanwhile, Adjoint is a dumb wrapper and will keep wrapping the wrapper if called multiple times.

5

Many thanks to both of you!

My limited understanding is that bar(sfoo) and bar(mfoo) are distinct by dispatch of the argument. Are they not?

Extending MMatrix{3,3} to MMatrix{3,3,Float64,9} does remove the type instability. I would like to beter understand why extending {3,3} to {3,3,Float64,9} is not required in the SMatrix case.

6

You need to add type annotations, as showed in my post above. Julia dispatches on types of arguments, not names of arguments.

As explained, if you actually tried to call bar(sfoo) the way you’ve written it, it would throw an error every time. This is why it looks “type stable”—it never returns and has no return type. Once you have a working implementation for both cases you will need a concrete field type in both cases, for example using parameters {3,3,Float64,9} (but see also the simpler and more generalizable solution at the bottom of my first post).

7

Thanks again.

Using

methods(bar)

I was able to gain a better understanding.

1 Like