fix: set-indexing into subarray

[avik-pal]

fixes #1636

[avik-pal]

module @reactant_test_sl... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<10x2xf64> {tf.aliasing_output = 0 : i32}) -> tensor<10x2xf64> {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<10xf64>
    %c = stablehlo.constant dense<2> : tensor<10x1xi64>
    %c_0 = stablehlo.constant dense<1> : tensor<10x1xi64>
    %c_1 = stablehlo.constant dense<[[1], [3], [5], [7], [9], [11], [13], [15], [17], [19]]> : tensor<10x1xi64>
    %c_2 = stablehlo.constant dense<[[2, 10], [2, 10], [2, 10], [2, 10], [2, 10], [2, 10], [2, 10], [2, 10], [2, 10], [2, 10]]> : tensor<10x2xi64>
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<10x2xf64>) -> tensor<2x10xf64>
    %1 = stablehlo.subtract %c_1, %c_0 : tensor<10x1xi64>
    %2 = stablehlo.divide %1, %c : tensor<10x1xi64>
    %3 = stablehlo.concatenate %1, %2, dim = 1 : (tensor<10x1xi64>, tensor<10x1xi64>) -> tensor<10x2xi64>
    %4 = stablehlo.remainder %3, %c_2 : tensor<10x2xi64>
    %5 = "stablehlo.scatter"(%0, %4, %cst) <{scatter_dimension_numbers = #stablehlo.scatter<inserted_window_dims = [0, 1], scatter_dims_to_operand_dims = [0, 1], index_vector_dim = 1>}> ({
    ^bb0(%arg1: tensor<f64>, %arg2: tensor<f64>):
      stablehlo.return %arg2 : tensor<f64>
    }) : (tensor<2x10xf64>, tensor<10x2xi64>, tensor<10xf64>) -> tensor<2x10xf64>
    %6 = stablehlo.transpose %5, dims = [1, 0] : (tensor<2x10xf64>) -> tensor<10x2xf64>
    return %6 : tensor<10x2xf64>
  }
}

we should be able to do this with a dus but currently the codegne is quite bad

[codecov]

Codecov Report

❌ Patch coverage is 90.00000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 42.60%. Comparing base (6bdd567) to head (a1279fa).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
src/TracedRArray.jl	50.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1639      +/-   ##
==========================================
+ Coverage   42.55%   42.60%   +0.04%     
==========================================
  Files         123      123              
  Lines       21901    21906       +5     
==========================================
+ Hits         9321     9333      +12     
+ Misses      12580    12573       -7     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.