New "easy rule" system for Enzyme

Specific Domains Machine Learning
, ,
1

Hi there @wsmoses and friends,

Thanks for the new easy rule system! Making rules more accessible is a great idea to get more people into autodiff (that’s how I got hooked in the first place).

I’m curious to understand how it works for functions with vector inputs and outputs, especially when their size may be large:

  • Does everything have to be stored in tuples of tuples? Will it strain the compiler if people (programmatically) define tuples with a million elements because their arrays have that size?
  • If the user only specifies the equivalent of a Jacobian matrix, how does Enzyme internally generate efficient JVPs and VJPs that don’t require materializing the full matrix?

Thanks again!

1 Like
2

The docstring has some more information

Does everything have to be stored in tuples of tuples? Will it strain the compiler if people (programmatically) define tuples with a million elements because their arrays have that size?

As far as I understand it no. The length of the tuple is bounded by the number of arguments.

As an example from the tests (cleaned up an actually using the correct jacobian)

# 3 -> 2 vector valued function
function f(v)
    return [sin(v[1]), cos(v[2]) * v[3]]
end

# Return the jacobian of f 
function jacobian_f(v)
   [cos(v[1]), 0, 0
   0 , -v[3]*sin(v[2]), cos(v[2])]
end

EnzymeRules.@easy_rule(
    f(v),
    @setup(),
    (jacobian_f(v),)
)
1 Like
3

Rules support arrays or scalars as inputs/outputs. There’s one row in the easyrule per each input to the function (either a scalar input or an array input). The corresponding rule expression then should be an abstractarray of shape (size(output[j])..., size(input[i])...) (the docs should say this, if not clear help us make it more clear :slight_smile: ). It needs to be an abstractarray, it need not be a materialized dense array (which can therefore avoid materialization).

1 Like
4

My bad, I had read the docstring it but intuitively I understood the notation

@easy_rule(f(x₁::T1, x₂::T2, ...),
             @setup(statement₁, statement₂, ...),
             (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
             (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
             ...)

as referring to a function

f(x1, x2, ...) = [f1, f2, ...]  # an array output

instead of what was probably meant, which is

f(x1, x2, ...) = (f1, f2, ...)  # a tuple output

Retrospectively, I think this confusion of mine stems from the idea that Julia functions don’t really have “several outputs”, they just output a tuple or an array which is in the end one single object. Not sure how it can be worded better though.

So if I want to implement an easy rule where the Jacobian is not materialized at each forward or reverse pass (which is n times slower than necessary), I have to define a new AbstractArray subtype which implements multiplication using a JVP and transpose-multiplication using a VJP? Something like this?

struct MyLazyJacobian <: AbstractArray
    # stuff
end

function Base.:*(J::MyLazyJacobian, v)
    #stuff
end

function Base.:*(J::Transpose{MyLazyJacobian}, v)
    #stuff
end
5

So if you wanted to have a custom materialization (rather than general array), I think currently you’d need:

Jacobian * Number 
LinearAlgebra.dot(Jacobian, Abstract Array of same dim)
LinearAlgebra.dot(Abstract Array of same dim, Jacobian)
Base.reshape
Base.length
Base.similar
LinearAlgebra.mul!(AbstractArray, Jacobian, AbstractArray)

Alternatively (and not presently well trodden since we’ve currently assumed the jacobian is a standard number/abstracyarray), you can override EnzymeCore.EnzymeRules.multiply_fwd_into(previous, jacobian, dx)
which should logically compute previous += jacobian * dx, if previous is not nothing, and return anew if it is nothing (and this by default will call into the utilities above, if not specialized).

1 Like