# How to differentiate functions involving arrays with one non-zero element in Zygote

**URL:** <https://discourse.julialang.org/t/how-to-differentiate-functions-involving-arrays-with-one-non-zero-element-in-zygote/127741>\
**Category:** General Usage\
**Tags:** zygote, autodiff\
**Created:** [April 5, 2025, 12:24pm UTC](https://discourse.julialang.org/t/how-to-differentiate-functions-involving-arrays-with-one-non-zero-element-in-zygote/127741 "2025-04-05T12:24:12Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![albertomercurio](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albertomercurio/32/27051_2.png) [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Post date:** [April 5, 2025, 12:24pm UTC](https://discourse.julialang.org/t/how-to-differentiate-functions-involving-arrays-with-one-non-zero-element-in-zygote/127741/1 "2025-04-05T12:24:12Z")

</div>

Hello,

I need to differentiate using Zygote.jl a function like

```julia
using LinearAlgebra

function my_fidelity2(p)
  ψg = zeros(ComplexF64, 2)
  ψg[2] = 1.0
  ψe = zeros(ComplexF64, 2)
  ψe[1] = 1.0

  ψgg = kron(ψg, ψg)
  ψee = kron(ψe, ψe)
  ψ = normalize(p[1] * ψgg + p[2] * ψee)

  return abs(dot(ψ, ψgg))
end

```

But here I have a mutating arrays, and so Zygote.jl complains. The differentiation works with Enzyme.jl, aside from the BLAS warning posted in [this issue](https://github.com/EnzymeAD/Enzyme.jl/issues/2349).

What should I do to make it working with Zygote.jl? I could for sure moving the definition of `ψg` and `ψe` outside the function. But what if I want to keep it inside it. Is there a way to generate an array with only one nonzero value that is supported by Zygote.jl?

I was thinking that I could create a function like

```julia
function my_state(N, i)
  ψ = zeros(ComplexF64, N)
  ψ[i] = 1.0
  return ψ
end

```

And the defining a custom `rrule` with ChainRulesCore.jl. I’m new to the autodiff world, so I have no idea of how to set it, or if this is the best solution.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [April 5, 2025, 12:36pm UTC](https://discourse.julialang.org/t/how-to-differentiate-functions-involving-arrays-with-one-non-zero-element-in-zygote/127741/2 "2025-04-05T12:36:36Z")

</div>

```julia
function my_fidelity2(p)
  ψg = zeros(ComplexF64, 1)
  ψg = [1.0; ψg]
  ψe = zeros(ComplexF64, 1)
  ψe = [1.0; ψe]

  ψgg = kron(ψg, ψg)
  ψee = kron(ψe, ψe)
  ψ = normalize(p[1] * ψgg + p[2] * ψee)

  return abs(dot(ψ, ψgg))
end

```

---

<div class="post-metadata">

**Author:** ![albertomercurio](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albertomercurio/32/27051_2.png) [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Post date:** [April 5, 2025, 12:40pm UTC](https://discourse.julialang.org/t/how-to-differentiate-functions-involving-arrays-with-one-non-zero-element-in-zygote/127741/3 "2025-04-05T12:40:31Z")

</div>

Thanks, very simple!

---

<div class="post-metadata">

**Author:** ![abraemer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/abraemer/32/51403_2.png) [@abraemer](https://discourse.julialang.org/u/abraemer)\
**Post date:** [April 7, 2025, 5:42am UTC](https://discourse.julialang.org/t/how-to-differentiate-functions-involving-arrays-with-one-non-zero-element-in-zygote/127741/4 "2025-04-07T05:42:17Z")

</div>

I think you could simplify this alot. Doesn’t it just calculate:

```julia
function my_fidelity3(p)
    return abs(p[1])/hypot(p[1], p[2])
end

```

---

<div class="post-metadata">

**Author:** ![albertomercurio](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albertomercurio/32/27051_2.png) [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Post date:** [April 13, 2025, 9:27am UTC](https://discourse.julialang.org/t/how-to-differentiate-functions-involving-arrays-with-one-non-zero-element-in-zygote/127741/5 "2025-04-13T09:27:15Z")

</div>

Yes, but I made this example just as a MWE. In principle the two vectors can be very general. Thanks anyway for this optimized solution.