Loosen index restriction

[wsmoses]

cc @mcabbott for testing with reactant

[wsmoses]

okay with the Reactant PR (which depends on this being tagged: EnzymeAD/Reactant.jl#1278), everything seems good:

(base) wmoses@hydra:~/git/Reactant.jl$ cat oh.jl 
using OneHotArrays, Reactant

function oh_mul(a, b)
  return a * b
end

m = onehotbatch([10, 20, 30, 10, 10], 10:10:40)
r_m = Reactant.to_rarray(m)
a = rand(100,4)
r_a = Reactant.to_rarray(a)
@show @code_hlo oh_mul(r_a, r_m)
r_res = @jit oh_mul(r_a, r_m)
res = oh_mul(a, m)

@show convert(Array, r_res) ≈ res

(base) wmoses@hydra:~/git/Reactant.jl$ julia --project=oh oh.jl 
Precompiling ReactantOneHotArraysExt...
  1 dependency successfully precompiled in 3 seconds. 105 already precompiled.
AssertionError("Could not find registered platform with name: \"cuda\". Available platform names are: ")
#= /home/wmoses/git/Reactant.jl/oh.jl:11 =# @code_hlo(oh_mul(r_a, r_m)) = module @reactant_oh_mul attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<4x100xf64>, %arg1: tensor<5xui32>) -> tensor<5x100xf64> {
    %c = stablehlo.constant dense<1> : tensor<5xi64>
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<4x100xf64>) -> tensor<100x4xf64>
    %1 = stablehlo.convert %arg1 : (tensor<5xui32>) -> tensor<5xi64>
    %2 = stablehlo.subtract %1, %c : tensor<5xi64>
    %3 = "stablehlo.gather"(%0, %2) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1], index_vector_dim = 1>, slice_sizes = array<i64: 100, 1>}> : (tensor<100x4xf64>, tensor<5xi64>) -> tensor<100x5xf64>
    %4 = stablehlo.transpose %3, dims = [1, 0] : (tensor<100x5xf64>) -> tensor<5x100xf64>
    return %4 : tensor<5x100xf64>
  }
}
convert(Array, r_res) ≈ res = true

[avik-pal]

cc @mcabbott for a review

[avik-pal]

@mcabbott bump on this