EnzymeAD / EnzymeAD/Reactant.jl
Unexpected output type when multiplying reshaped ConcreteRArrays
- Dominant language
- Julia
- Stars
- 370
- Forks
- 74
- Avg merge
- 18h 47m
- Merged PRs (30d)
- 30
Description
### Description
When reshaping a `ConcreteRArray` and performing pairwise multiplication, the result returns a standard `Array` rather than retaining the `ConcreteRArray` type. This behavior seems to deviate from the expected functionality, as it would be ideal for the result of such operations to preserve the `ConcreteRArray` type to ensure consistency in further operations. This issue becomes problematic when used within `Lux` neural networks. Specifically, when performing matrix operations on reshaped `ConcreteRArray` objects, the type mismatch between `ConcreteRArray` and `Array` triggers an `ArgumentError` due to objects being on different devices (`ReactantDevice` and `CPUDevice`). This disrupts the execution of the model and leads to errors in subsequent operations since `Lux.Utils.make_abstract_matrix` is involved internally.
### Steps to reproduce
```julia
Reactant.set_default_backend("cpu")
const xdev = reactant_device()
x = rand(Float32, 2, 10, 3) |> xdev #ConcreteRArray{Float32, 3, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}
reshape(x, 2, :) .* reshape(x, 2, :) #Matrix{Float32}
```
```julia
model = Chain(Parallel(.*, Dense(2, 16, tanh), Dense(2, 16, tanh)), Dense(16, 1))
ps, st = Lux.setup(Random.default_rng(), model) |> xdev
model(x, ps, st)
```
This triggers an error:
```julia
ERROR: ArgumentError: Objects are on devices with different types: ReactantDevice and CPUDevice.
Stacktrace:
[1] combine_devices(T1::Type{ReactantDevice}, T2::Type{CPUDevice})
@ MLDataDevices.Internal ~/.julia/packages/MLDataDevices/uhCbD/src/internal.jl:125
[2] macro expansion
@ ~/.julia/packages/MLDataDevices/uhCbD/src/internal.jl:205 [inlined]
[3] unrolled_mapreduce
@ ~/.julia/packages/MLDataDevices/uhCbD/src/internal.jl:192 [inlined]
[4] unrolled_mapreduce(f::typeof(get_device_type), op::typeof(MLDataDevices.Internal.combine_devices), itr::Tuple{ConcreteRArray{…}, Matrix{…}, ConcreteRArray{…}})
@ MLDataDevices.Internal ~/.julia/packages/MLDataDevices/uhCbD/src/internal.jl:183
[5] get_device_type(x::Tuple{ConcreteRArray{Float32, 2, 1, Reactant.Sharding.ShardInfo{…}}, Matrix{Float32}, ConcreteRArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{…}}})
@ MLDataDevices.Internal ~/.julia/packages/MLDataDevices/uhCbD/src/internal.jl:160
[6] get_device_type(x::Tuple{ConcreteRArray{Float32, 2, 1, Reactant.Sharding.ShardInfo{…}}, Matrix{Float32}, ConcreteRArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{…}}})
@ MLDataDevices ~/.julia/packages/MLDataDevices/uhCbD/src/public.jl:370
[7] internal_operation_mode(xs::Tuple{ConcreteRArray{Float32, 2, 1, Reactant.Sharding.ShardInfo{…}}, Matrix{Float32}, ConcreteRArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{…}}})
@ LuxLib ~/.julia/packages/LuxLib/kH9PB/src/traits.jl:217
[8] select_fastest_activation(::typeof(identity), ::ConcreteRArray{Float32, 2, 1, Reactant.Sharding.ShardInfo{…}}, ::Matrix{Float32}, ::ConcreteRArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{…}})
@ LuxLib.Impl ~/.julia/packages/LuxLib/kH9PB/src/impl/activation.jl:129
[9] fused_dense_bias_activation(σ::typeof(identity), weight::ConcreteRArray{Float32, 2, 1, Reactant.Sharding.ShardInfo{…}}, x::Matrix{Float32}, b::ConcreteRArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{…}})
@ LuxLib.API ~/.julia/packages/LuxLib/kH9PB/src/api/dense.jl:35
[10] (::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True})(x::Array{Float32, 3}, ps::@NamedTuple{weight::ConcreteRArray{…}, bias::ConcreteRArray{…}}, st::@NamedTuple{})
@ Lux ~/.julia/packages/Lux/TbS8R/src/layers/basic.jl:343
[11] apply
@ ~/.julia/packages/LuxCore/8mVob/src/LuxCore.jl:155 [inlined]
[12] macro expansion
@ ~/.julia/packages/Lux/TbS8R/src/layers/containers.jl:0 [inlined]
[13] applychain(layers::@NamedTuple{…}, x::ConcreteRArray{…}, ps::@NamedTuple{…}, st::@NamedTuple{…})
@ Lux ~/.julia/packages/Lux/TbS8R/src/layers/containers.jl:482
[14] (::Chain{@NamedTuple{…}, Nothing})(x::ConcreteRArray{Float32, 3, 1, Reactant.Sharding.ShardInfo{…}}, ps::@NamedTuple{layer_1::@NamedTuple{…}, layer_2::@NamedTuple{…}}, st::@NamedTuple{layer_1::@NamedTuple{…}, layer_2::@NamedTuple{}})
@ Lux ~/.julia/packages/Lux/TbS8R/src/layers/containers.jl:480
[15] top-level scope
@ REPL[47]:1
Some type information was truncated. Use `show(err)` to see complete types.
```
```julia
[b2108857] Lux v1.7.0
[3c362404] Reactant v0.2.31
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.