# Split an array into multiple smaller arrays - alternative to unsafe_wrap without views

I’m wondering if this is a good use case for unsafe_wrap, or if there is a better way to do this. Basically I have a vector, and I want to be able to interpret the vector as a set of arrays. I cannot use subarrays or similar, as something about them doesn’t play nice with Zygote/CuArrays (https://github.com/FluxML/Zygote.jl/issues/600).

``````using Revise, Flux, Zygote

struct mwestruct{TΠ, TA}
Π::TΠ
A::TA
B::TA
end

#constructor
function mwestruct(K::Int, A::TA, B::TA) where TA<:AbstractMatrix

i = (i->(i-1)%K+1).(1:K^2)
Π = vcat(vec(A), vec(B)) #this is the parameter vector

startA = 1
p = pointer(Π, startA)

startB = length(vec(A))+1
p = pointer(Π, startB)

end

#lets call this the loss function (much simpler than the non-mwe version)
function (m::mwestruct)(x)
sum(abs.((m.A .-  m.A.*m.B .- diag(m.B))  ./ (x .+ 1f0)))
end

Flux.@functor mwestruct
Flux.trainable(a::mwestruct) = (a.A,a.B)

#this computes the gradient as a vector corresponding to the parameter vector

function zygotemwe(N)
A = Matrix{Float32}(reshape(collect(1:N^2), (N,N)))
B = Matrix{Float32}(reshape(collect(1:N^2), (N,N)))
x = Matrix{Float32}(reshape(collect(1:N^2), (N,N))) .* 2

s=mwestruct(N, A, B)