# Replacing static arrays for use with Enzyme.jl

**URL:** <https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278>\
**Category:** General Usage\
**Tags:** question, staticarrays, enzyme\
**Created:** [January 25, 2024, 9:34pm UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278 "2024-01-25T21:34:11Z")\
**Posts on this page:** 11\
**Page:** 1

<div class="post-metadata">

**Author:** ![jwalker](https://avatars.discourse-cdn.com/v4/letter/j/9fc29f/32.png) [@jwalker](https://discourse.julialang.org/u/jwalker)\
**Post date:** [January 25, 2024, 9:34pm UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278/1 "2024-01-25T21:34:12Z")

</div>

I’m implementing AD in a code that makes use of types built upon `SLArray` and `SArray`, and am trying to understand the best way toward converting things so that the code plays well with Enzyme.

From my understanding, to work with Enzyme, 1) all arrays and custom data types must be mutable, and 2) functions must mutate the arrays in-place and not return them. Is this correct?

If so, then my my naive strategy would be to replace `SArray` with `MArray` and to replace `SLArray` with `LArray`. Does this seem like a good route?

If this thinking is correct, then I have some general questions. Is there an unavoidable performance trade-off between using `SArray` versus `MArray` in order to use Enzyme? Also, I noticed in LabelledArrays.jl that `SLArray` is based on `SArray` but that `LArray` is not based on `MArray` as I might expect. Is this the intended end state, or is it just because no one has implemented it yet? Maybe I’m misunderstanding this.

---

<div class="post-metadata">

**Author:** ![jwalker](https://avatars.discourse-cdn.com/v4/letter/j/9fc29f/32.png) [@jwalker](https://discourse.julialang.org/u/jwalker)\
**Post date:** [January 25, 2024, 10:30pm UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278/2 "2024-01-25T22:30:16Z")

</div>

To possibly answer that last question, I guess it’s that `LArray` has the flexibility to allow an `<: AbstractArray` type, so that something like `MArray` would still work.

---

<div class="post-metadata">

**Author:** ![jgreener64](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jgreener64/32/2483_2.png) [@jgreener64](https://discourse.julialang.org/u/jgreener64)\
**Post date:** [January 26, 2024, 12:06am UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278/3 "2024-01-26T00:06:19Z")

</div>

Enzyme generally works well with static arrays. For example:

```julia
julia> using Enzyme, StaticArrays

julia> x = rand(SVector{3, Float64}, 10)
10-element Vector{SVector{3, Float64}}:
 [0.703458204835627, 0.7121617637366355, 0.7313927507314387]
 [0.4621987150741128, 0.2154237571455806, 0.10735184165752176]
 [0.545521087048848, 0.2192157907272828, 0.7063366942886221]
 [0.29922059211540786, 0.00022592689698108792, 0.4566363080401129]
 [0.0804002475278649, 0.7924864781611491, 0.7757979667815256]
 [0.4175042850253522, 0.6422189150032824, 0.3510186112988213]
 [0.7120487718195188, 0.09890390308324437, 0.7345662304826344]
 [0.08079153987374355, 0.4615163823192098, 0.23119102793791468]
 [0.6539763154678109, 0.6466826908349046, 0.8395969123227822]
 [0.3203653937684048, 0.6881807184735692, 0.12676012258252278]

julia> dx = zero(x)
10-element Vector{SVector{3, Float64}}:
 [0.0, 0.0, 0.0]
 [0.0, 0.0, 0.0]
 [0.0, 0.0, 0.0]
 [0.0, 0.0, 0.0]
 [0.0, 0.0, 0.0]
 [0.0, 0.0, 0.0]
 [0.0, 0.0, 0.0]
 [0.0, 0.0, 0.0]
 [0.0, 0.0, 0.0]
 [0.0, 0.0, 0.0]

julia> function f(x, y)
           sum(sum(y .* x))
       end
f (generic function with 1 method)

julia> f(x, 2.0)
27.626299890124855

julia> autodiff(Reverse, f, Active, Duplicated(x, dx), Active(2.0))
((nothing, 13.813149945062424),)

julia> dx
10-element Vector{SVector{3, Float64}}:
 [2.0, 2.0, 2.0]
 [2.0, 2.0, 2.0]
 [2.0, 2.0, 2.0]
 [2.0, 2.0, 2.0]
 [2.0, 2.0, 2.0]
 [2.0, 2.0, 2.0]
 [2.0, 2.0, 2.0]
 [2.0, 2.0, 2.0]
 [2.0, 2.0, 2.0]
 [2.0, 2.0, 2.0]

julia> g(x) = 2 * sum(x)
g (generic function with 1 method)

julia> autodiff(Reverse, g, Active, Active(SVector(1.0, 2.0, 3.0)))
(([2.0, 2.0, 2.0],),)

```

Maybe give it a go as-is and see if you get any errors.

---

<div class="post-metadata">

**Author:** ![jwalker](https://avatars.discourse-cdn.com/v4/letter/j/9fc29f/32.png) [@jwalker](https://discourse.julialang.org/u/jwalker)\
**Post date:** [January 26, 2024, 12:54am UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278/4 "2024-01-26T00:54:45Z")

</div>

Huh. I did try as-is before but got this error, which lead me to believe that SArray was the issue (this is partial):

```julia
  Got exception outside of a @test
  setfield!: immutable struct of type SArray cannot be changed
  Stacktrace:
    [1] rt_jl_getfield_rev(::SMatrix{3, 3, Float64, 9}, ::Base.RefValue{NTuple{9, Float64}}, ::Type{Val{:data}}, ::Val{false})
      @ Enzyme.Compiler ~/.julia/packages/Enzyme/Dd2LU/src/rules/typeunstablerules.jl:257
    [2] getindex
      @ ~/.julia/packages/StaticArrays/eGKzB/src/SArray.jl:62 [inlined]
    [3] view
      @ ~/.julia/packages/StaticArrays/eGKzB/src/abstractarray.jl:291 [inlined]
    [4] getindex

```

This error remains, even after I make the custom type a `mutable struct` rather than `struct`

---

<div class="post-metadata">

**Author:** ![vchuravy](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/vchuravy/32/8_2.png) [@vchuravy](https://discourse.julialang.org/u/vchuravy)\
**Post date:** [January 26, 2024, 1:05am UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278/5 "2024-01-26T01:05:41Z")

</div>

That seems to be a bug with one of the custom rules, so if you have a MWE that triggers that please open an issue.

---

<div class="post-metadata">

**Author:** ![jwalker](https://avatars.discourse-cdn.com/v4/letter/j/9fc29f/32.png) [@jwalker](https://discourse.julialang.org/u/jwalker)\
**Post date:** [January 31, 2024, 4:17pm UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278/6 "2024-01-31T16:17:02Z")

</div>

Here’s an MWE. I haven’t opened an issue because I’m not sure whether it’s on my end or not. The issue seems to be the `getindex` function.

```julia
using Enzyme
using Random 
using StaticArrays

abstract type AbstractBasisType end

struct Contravariant <: AbstractBasisType end

struct CurvilinearBasisVectors{N, T, C, B, V <: AbstractBasisType} <: StaticMatrix{N, N, T}
    __x::Union{SMatrix{N, N, T}}
    function CurvilinearBasisVectors{N, T, C, B, V}(b::AbstractMatrix) where {N, T, C, B, V}
        return new{N, T, C, B, V}(SMatrix{N, N, T}(b))
    end
end

Base.@propagate_inbounds function Base.getindex(v::CurvilinearBasisVectors{N, T, C, B, V}, i::Int) where {N, T, C, B, V}
    return view(getfield(v, :__x), i)[]
end

basis_labels = (:∇x, :∇y, :∇z);
coord_labels = (:x, :y, :z);

a = CurvilinearBasisVectors{3,Float64,basis_labels,coord_labels,Contravariant}(
    rand(3,3),
);
da = CurvilinearBasisVectors{3,Float64,basis_labels,coord_labels,Contravariant}(
    zeros(3,3),
);

@show a,da

function f(x, y)
    sum(sum(y .* x))
end

@show f(a, 2.0)

@show autodiff(Reverse, f, Active, Duplicated(a, da), Active(2.0))
@show da

```

---

<div class="post-metadata">

**Author:** ![vchuravy](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/vchuravy/32/8_2.png) [@vchuravy](https://discourse.julialang.org/u/vchuravy)\
**Post date:** [January 31, 2024, 4:33pm UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278/7 "2024-01-31T16:33:54Z")

</div>

For duplicated in reverse mode the arguments need to be `mutable` you defined `CurvilinearBasisVectors` as a `struct`, if you use `mutable struct` it should work.

---

<div class="post-metadata">

**Author:** ![jwalker](https://avatars.discourse-cdn.com/v4/letter/j/9fc29f/32.png) [@jwalker](https://discourse.julialang.org/u/jwalker)\
**Post date:** [January 31, 2024, 4:39pm UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278/8 "2024-01-31T16:39:44Z")

</div>

That was the first thing I tried (but forgot to mention it in the OP). Trying it again with this MWE, I get the same error:

```julia
ERROR: LoadError: setfield!: immutable struct of type SArray cannot be changed
Stacktrace:
  [1] rt_jl_getfield_rev(::SMatrix{3, 3, Float64, 9}, ::Base.RefValue{NTuple{9, Float64}}, ::Type{Val{:data}}, ::Val{false})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/rules/typeunstablerules.jl:257
  [2] getindex
    @ ~/.julia/packages/StaticArrays/eGKzB/src/SArray.jl:62 [inlined]
  [3] view
    @ ~/.julia/packages/StaticArrays/eGKzB/src/abstractarray.jl:291 [inlined]
  [4] getindex
    @ ~/software/julia/enzyme/test-basis.jl:17 [inlined]
  [5] macro expansion
    @ ~/.julia/packages/StaticArrays/eGKzB/src/broadcast.jl:135 [inlined]
  [6] __broadcast
    @ ~/.julia/packages/StaticArrays/eGKzB/src/broadcast.jl:123 [inlined]
  [7] _broadcast
    @ ~/.julia/packages/StaticArrays/eGKzB/src/broadcast.jl:119 [inlined]
  [8] copy
    @ ~/.julia/packages/StaticArrays/eGKzB/src/broadcast.jl:60 [inlined]
  [9] materialize
    @ ./broadcast.jl:903 [inlined]
 [10] f
    @ ~/software/julia/enzyme/test-basis.jl:33 [inlined]
 [11] diffejulia_f_2626wrap
    @ ~/software/julia/enzyme/test-basis.jl:0
 [12] macro expansion
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/compiler.jl:5306 [inlined]
 [13] enzyme_call(::Val{…}, ::Ptr{…}, ::Type{…}, ::Type{…}, ::Val{…}, ::Type{…}, ::Type{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Active{…}, ::Float64)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/compiler.jl:4984
 [14] (::Enzyme.Compiler.CombinedAdjointThunk{…})(::Const{…}, ::Duplicated{…}, ::Vararg{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/compiler.jl:4926
 [15] autodiff(::ReverseMode{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Vararg{…})
    @ Enzyme ~/.julia/packages/Enzyme/jOGYG/src/Enzyme.jl:215
 [16] autodiff(::ReverseMode{false, FFIABI}, ::typeof(f), ::Type, ::Duplicated{CurvilinearBasisVectors{…}}, ::Vararg{Any})
    @ Enzyme ~/.julia/packages/Enzyme/jOGYG/src/Enzyme.jl:224
 [17] macro expansion
    @ show.jl:1181 [inlined]
 [18] top-level scope
    @ ~/software/julia/enzyme/test-basis.jl:38
 [19] include(fname::String)
    @ Base.MainInclude ./client.jl:489
 [20] top-level scope
    @ REPL[2]:1

```

---

<div class="post-metadata">

**Author:** ![jwalker](https://avatars.discourse-cdn.com/v4/letter/j/9fc29f/32.png) [@jwalker](https://discourse.julialang.org/u/jwalker)\
**Post date:** [January 31, 2024, 4:44pm UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278/9 "2024-01-31T16:44:50Z")

</div>

If I also change the `SMatrix` in the example above to `MMatrix`, then I get a different error:

```julia
ERROR: LoadError: Enzyme Mutability Error: Cannot add one in place to immutable value (Base.RefValue{Float64}(0.0), Base.RefValue{Float64}(0.0), Base.RefValue{Float64}(0.0), Base.RefValue{Float64}(0.0), Base.RefValue{Float64}(0.0), Base.RefValue{Float64}(0.0), Base.RefValue{Float64}(0.0), Base.RefValue{Float64}(0.0), Base.RefValue{Float64}(0.0))
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] runtime_generic_rev(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, tape::Enzyme.Compiler.Tape{…}, f::Type{…}, df::Type{…}, primal_1::NTuple{…}, shadow_1_1::NTuple{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/rules/jitrules.jl:242
  [3] _broadcast
    @ ~/.julia/packages/StaticArrays/eGKzB/src/broadcast.jl:120 [inlined]
  [4] copy
    @ ~/.julia/packages/StaticArrays/eGKzB/src/broadcast.jl:60 [inlined]
  [5] materialize
    @ ./broadcast.jl:903 [inlined]
  [6] f
    @ ~/software/julia/enzyme/test-basis.jl:33 [inlined]
  [7] diffejulia_f_2725wrap
    @ ~/software/julia/enzyme/test-basis.jl:0
  [8] macro expansion
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/compiler.jl:5306 [inlined]
  [9] enzyme_call(::Val{…}, ::Ptr{…}, ::Type{…}, ::Type{…}, ::Val{…}, ::Type{…}, ::Type{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Active{…}, ::@NamedTuple{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/compiler.jl:4984
 [10] (::Enzyme.Compiler.AdjointThunk{…})(::Const{…}, ::Duplicated{…}, ::Vararg{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/compiler.jl:4932
 [11] autodiff(::ReverseMode{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Vararg{…})
    @ Enzyme ~/.julia/packages/Enzyme/jOGYG/src/Enzyme.jl:203
 [12] autodiff(::ReverseMode{false, FFIABI}, ::typeof(f), ::Type, ::Duplicated{CurvilinearBasisVectors{…}}, ::Vararg{Any})
    @ Enzyme ~/.julia/packages/Enzyme/jOGYG/src/Enzyme.jl:224
 [13] macro expansion
    @ show.jl:1181 [inlined]
 [14] top-level scope
    @ ~/software/julia/enzyme/test-basis.jl:38
 [15] include(fname::String)
    @ Base.MainInclude ./client.jl:489
 [16] top-level scope
    @ REPL[2]:1

```

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [January 31, 2024, 9:15pm UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278/10 "2024-01-31T21:15:36Z")

</div>

So this specifically is a “not-yet implemented” for a type unstable update of an immutable variable.

We should add support for this, but in the interim, can you type stabilize your code?

Regardless please open an issue.

---

<div class="post-metadata">

**Author:** ![jwalker](https://avatars.discourse-cdn.com/v4/letter/j/9fc29f/32.png) [@jwalker](https://discourse.julialang.org/u/jwalker)\
**Post date:** [January 31, 2024, 11:04pm UTC](https://discourse.julialang.org/t/replacing-static-arrays-for-use-with-enzyme-jl/109278/11 "2024-01-31T23:04:30Z")

</div>

Issue opened. [https://github.com/EnzymeAD/Enzyme.jl/issues/1263](https://github.com/EnzymeAD/Enzyme.jl/issues/1263)

And I’ll try to type stabilize the code. Thanks for the suggestion.