Handling random numbers correctly

[avik-pal]

julia> function test_rand(x)
           y = similar(x)
           randn!(y)
           return y
       end
test_rand (generic function with 1 method)

julia> test_rand(rand(2, 3))
2×3 Matrix{Float64}:
 -0.856593   2.09777  -1.47022
 -0.587353  -1.63428   0.101531

julia> test_rand(rand(2, 3))
2×3 Matrix{Float64}:
 -0.169647  0.20255    1.73269
  1.99007   0.0300865  1.00466

julia> @code_hlo test_rand(Reactant.to_rarray(rand(2, 3)))
module {
  func.func @main(%arg0: tensor<3x2xf64>) -> tensor<3x2xf64> {
    %cst = stablehlo.constant dense<[[-0.14388180121311803, -1.7968975324220513], [0.40886485616904072, -0.26052122391271643], [-2.1956924346267952, -0.047720339420302453]]> : tensor<3x2xf64>
    return %cst : tensor<3x2xf64>
  }
}

julia> @code_hlo test_rand(Reactant.to_rarray(rand(2, 3)))
module {
  func.func @main(%arg0: tensor<3x2xf64>) -> tensor<3x2xf64> {
    %cst = stablehlo.constant dense<[[-0.39078327584480177, 0.46390776775279169], [-0.2758724528591438, -2.4972617545675377], [0.51502111596476896, -0.56535270904132628]]> : tensor<3x2xf64>
    return %cst : tensor<3x2xf64>
  }
}

julia> fn = @compile test_rand(Reactant.to_rarray(rand(2, 3)))
Reactant.Compiler.Thunk{Symbol("##test_rand_reactant#655815")}()

julia> fn(Reactant.to_rarray(rand(2, 3)))
2×3 ConcreteRArray{Float64, 2}:
 -0.788288  -1.67923   0.00359754
  0.268338   1.13756  -0.907077

julia> fn(Reactant.to_rarray(rand(2, 3)))
2×3 ConcreteRArray{Float64, 2}:
 -0.788288  -1.67923   0.00359754
  0.268338   1.13756  -0.907077

This is not unexpected from how we do the tracing but definitely not the preferred behavior

[avik-pal]

https://openxla.org/stablehlo/spec#rng

[wsmoses]

so ideally we should be able to use that and make a traced rng, I think?

[wsmoses]

also once my absint overload stuff gets in from the cuda side, we probably should override random.default_rng or w/e

[avik-pal]

so ideally we should be able to use that and make a traced rng, I think?

We need the following:

1. rand(/n)!(r::TracedRNG, ....)
2. rand(/n)!(r::AbstractRNG, ::TracedRArray, ....)

There is also a direct Sampler API I think but don't recall how to overload that correctly

[avik-pal]

Also noticed rng is being deprecated, we need to use https://openxla.org/stablehlo/spec#rng_bit_generator instead

[avik-pal]

Reference for implementation using the bit generator: https://github.com/JuliaRandom/Random123.jl/blob/master/src/common.jl#L125