# Generic array code that works with CUDA and Zygote?

**URL:** https://discourse.julialang.org/t/generic-array-code-that-works-with-cuda-and-zygote/104295
**Category:** GPU
**Tags:** cuda, zygote, cuarrays
**Created:** [September 27, 2023, 7:22am UTC](https://discourse.julialang.org/t/generic-array-code-that-works-with-cuda-and-zygote/104295 "2023-09-27T07:22:22Z")
**Posts on this page:** 1
**Page:** 1

<div class="post-metadata">

### Author: ![mchitre](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mchitre/32/13756_2.png) [@mchitre](https://discourse.julialang.org/u/mchitre)
#### Post date: [September 27, 2023, 7:22am UTC](https://discourse.julialang.org/t/generic-array-code-that-works-with-cuda-and-zygote/104295/1 "2023-09-27T07:22:22Z")

</div>

I am interested in writing a generic function that operates on an `AbstractArray` and produces some output. In doing so, it needs to create an intermediate array. The function should work with CPU arrays and GPU arrays, and the function needs to be differentiable with Zygote (so no mutation allowed).

Let me illustrate with an example:

```julia
function f1(x::AbstractArray)
  y = map(x -> x^2 + 1, eachindex(x))
  sum(x .* y)
end

```

This will work with most arrays and with Zygote, but will fail if `x` is a GPU array, as `y` won’t be on GPU and the `x .* y` can’t deal with a mix of on-GPU and on-CPU arrays.

This version:

```julia
function f2(x::AbstractArray)
  y = similar(x)
  map!(x -> x^2 + 1, y, eachindex(x))
  sum(x .* y)
end

```

will work with `x` being a GPU or a CPU array just fine, but will fail with Zygote due to mutation.

Is there a good idiomatic way to do write this generically without mutation?

P.S. This particular example could have been solved with `dot()`, but I’m asking a more general question about computation where I need `y` to be similar to `x` for computation in terms of the location (CPU/GPU) of the array.