fix: make eltype of Traced/Concrete Arrays to be respective RNumbers

[avik-pal]

No description provided.

[wsmoses]

I think we can get away with just doing this for the traced version and not the concrete variant (since a concrete number has higher overhead than otherwise, and it’ll obviously get optimized out for the traced case).

[avik-pal]

Lux tests will fail for now. Needs LuxDL/Lux.jl#1140

[mofeing]

what about #425 (comment)? isn't this way gonna be problematic too?

the case that comes to my mind is Yao. if the following example works, then it's fine from my side (actually, maybe i should add a test for this? i think i forgot to add them)

using YaoBlocks

f(x) = mat(Rx(x))
x = ConcreteRNumber(0.0)
@jit f(x)

EDIT: added #431 to test that this PR doesn't break integration with YaoBlocks

[avik-pal]

what about #425 (comment)? isn't this way gonna be problematic too?

the case that comes to my mind is Yao. if the following example works, then it's fine from my side (actually, maybe i should add a test for this? i think i forgot to add them)

using YaoBlocks

f(x) = mat(Rx(x))
x = ConcreteRNumber(0.0)
@jit f(x)

EDIT: added #431 to test that this PR doesn't break integration with YaoBlocks

julia> @code_hlo f(x)
module {
  func.func @main(%arg0: tensor<f64>) -> tensor<2x2xcomplex<f64>> {
    %c = stablehlo.constant dense<1> : tensor<i64>
    %c_0 = stablehlo.constant dense<0> : tensor<i64>
    %cst = stablehlo.constant dense<(0.000000e+00,-1.000000e+00)> : tensor<complex<f64>>
    %cst_1 = stablehlo.constant dense<2.000000e+00> : tensor<f64>
    %cst_2 = stablehlo.constant dense<(0.000000e+00,0.000000e+00)> : tensor<2x2xcomplex<f64>>
    %0 = stablehlo.divide %arg0, %cst_1 : tensor<f64>
    %1 = stablehlo.cosine %0 : tensor<f64>
    %2 = stablehlo.sine %0 : tensor<f64>
    %3 = stablehlo.convert %2 : (tensor<f64>) -> tensor<complex<f64>>
    %4 = stablehlo.multiply %cst, %3 : tensor<complex<f64>>
    %5 = stablehlo.convert %1 : (tensor<f64>) -> tensor<complex<f64>>
    %6 = stablehlo.reshape %5 : (tensor<complex<f64>>) -> tensor<1x1xcomplex<f64>>
    %7 = stablehlo.dynamic_update_slice %cst_2, %6, %c_0, %c_0 : (tensor<2x2xcomplex<f64>>, tensor<1x1xcomplex<f64>>, tensor<i64>, tensor<i64>) -> tensor<2x2xcomplex<f64>>
    %8 = stablehlo.dynamic_update_slice %7, %6, %c, %c : (tensor<2x2xcomplex<f64>>, tensor<1x1xcomplex<f64>>, tensor<i64>, tensor<i64>) -> tensor<2x2xcomplex<f64>>
    %9 = stablehlo.reshape %4 : (tensor<complex<f64>>) -> tensor<1x1xcomplex<f64>>
    %10 = stablehlo.dynamic_update_slice %8, %9, %c_0, %c : (tensor<2x2xcomplex<f64>>, tensor<1x1xcomplex<f64>>, tensor<i64>, tensor<i64>) -> tensor<2x2xcomplex<f64>>
    %11 = stablehlo.dynamic_update_slice %10, %9, %c, %c_0 : (tensor<2x2xcomplex<f64>>, tensor<1x1xcomplex<f64>>, tensor<i64>, tensor<i64>) -> tensor<2x2xcomplex<f64>>
    %12 = stablehlo.transpose %11, dims = [1, 0] : (tensor<2x2xcomplex<f64>>) -> tensor<2x2xcomplex<f64>>
    return %12 : tensor<2x2xcomplex<f64>>
  }
}

julia> @jit f(x)
2024-12-28 19:42:57.105196: I external/xla/xla/service/llvm_ir/llvm_command_line_options.cc:50] XLA (re)initializing LLVM with options fingerprint: 1154763464302348808
2×2 ConcreteRArray{ComplexF64, 2}:
 1.0+0.0im  0.0+0.0im
 0.0+0.0im  1.0+0.0im

[mofeing]

Great, then all in!

[wsmoses]

@avik-pal, since @mofeing's test case seems to pass (and also strictly this is an improvement to cases where the assumption is made in julia programs in that typeof(getindex(A, idx)) == eltype(A) is now actually satisfied by tracedrarrays), this is fine to merge when ready imo