feat: more JLL changes for #788

[avik-pal]

No description provided.

[giordano]

Test failures on this branch https://github.com/EnzymeAD/Reactant.jl/actions/runs/13470969798/job/37644418594#step:9:355

ArgumentError: unsafe_wrap: pointer 0x1 is not properly aligned to 8 bytes
  Stacktrace:
    [1] #unsafe_wrap#106
      @ ./pointer.jl:92 [inlined]
    [2] unsafe_wrap
      @ ./pointer.jl:90 [inlined]
    [3] convert(::Type{Reactant.XLA.OpSharding}, sharding::Reactant.XLA.JLOpSharding)
      @ Reactant.XLA ~/work/Reactant.jl/Reactant.jl/src/xla/Sharding.jl:90
    [4] Fix1
      @ ./operators.jl:1118 [inlined]
    [5] call_composed
      @ ./operators.jl:1044 [inlined]
    [6] ComposedFunction
      @ ./operators.jl:1041 [inlined]
    [7] iterate
      @ ./generator.jl:47 [inlined]
    [8] _collect(c::Vector{Base.RefValue{Reactant.XLA.JLOpSharding}}, itr::Base.Generator{Vector{Base.RefValue{Reactant.XLA.JLOpSharding}}, ComposedFunction{Base.Fix1{typeof(convert), Type{Reactant.XLA.OpSharding}}, typeof(getindex)}}, ::Base.EltypeUnknown, isz::Base.HasShape{1})
      @ Base ./array.jl:854
    [9] collect_similar
      @ ./array.jl:763 [inlined]
   [10] map
      @ ./abstractarray.jl:3286 [inlined]
   [11] get_parameter_shardings(exec::Reactant.XLA.PJRT.LoadedExecutable)
      @ Reactant.XLA.PJRT ~/work/Reactant.jl/Reactant.jl/src/xla/PJRT/LoadedExecutable.jl:61
   [12] convert(::Type{Reactant.XLA.OpSharding}, sharding_and_shape::Reactant.Sharding.ShardingWithShape{Reactant.Sharding.NamedSharding{2, 2, Tuple{Symbol, Nothing, Symbol}, 3}, Tuple{Int64, Int64, Int64}})
      @ Reactant.Sharding ~/work/Reactant.jl/Reactant.jl/src/Sharding.jl:181
   [13] convert
      @ ~/work/Reactant.jl/Reactant.jl/src/Sharding.jl:171 [inlined]
   [14] (::Reactant.Sharding.NamedSharding{2, 2, Tuple{Symbol, Nothing, Symbol}, 3})(client::Reactant.XLA.PJRT.Client, ::Nothing, x::Array{Float64, 3})
      @ Reactant.Sharding ~/work/Reactant.jl/Reactant.jl/src/Sharding.jl:107
   [15] ConcreteRArray(data::Array{Float64, 3}; client::Reactant.XLA.PJRT.Client, idx::Nothing, device::Nothing, sharding::Reactant.Sharding.NamedSharding{2, 2, Tuple{Symbol, Nothing, Symbol}, 3})
      @ Reactant ~/work/Reactant.jl/Reactant.jl/src/Types.jl:136
   [16] to_rarray_internal(x::Array{<:Union{Complex{Int16}, Complex{Int32}, Complex{Int64}, Complex{Int8}, Complex{UInt16}, Complex{UInt32}, Complex{UInt64}, Complex{UInt8}, Bool, Float16, Float32, Float64, Int16, Int32, Int64, Int8, UInt16, UInt32, UInt64, UInt8, Reactant.F8E4M3B11FNUZ, Reactant.F8E4M3FN, Reactant.F8E4M3FNUZ, Reactant.F8E5M2, Reactant.F8E5M2FNUZ, ComplexF64, ComplexF32}}, track_numbers::Type, sharding::Any)
      @ Reactant ~/work/Reactant.jl/Reactant.jl/src/Tracing.jl:1445
   [17] #to_rarray#89
      @ ~/work/Reactant.jl/Reactant.jl/src/Tracing.jl:1407 [inlined]

ere already introduced in #789 (but CI didn't finish running on that PR): https://github.com/EnzymeAD/Reactant.jl/actions/workflows/CI-localjll.yml?query=branch%3Amain

Is this because this needs adapting also on the Julia side (#788)?

[avik-pal]

Yeah. We need the julia changes as well for things to pass