more scatter-gather patterns

```
module @jit_sim attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
func.func @main() -> (tensor<1x402xf32> {jax.result_info = "result"}) {
%cst = stablehlo.constant dense<1.000000e+03> : tensor<1xf32>
%cst_0 = stablehlo.constant dense<999999.937> : tensor<1xf32>
%cst_1 = stablehlo.constant dense<2.500000e-02> : tensor<1xf32>
%cst_2 = stablehlo.constant dense<999.999938> : tensor<1xf32>
%cst_3 = stablehlo.constant dense<3.000000e-04> : tensor<2x1xf32>
%cst_4 = stablehlo.constant dense<-5.430000e+01> : tensor<2x1xf32>
%cst_5 = stablehlo.constant dense<-7.700000e+01> : tensor<2x1xf32>
%cst_6 = stablehlo.constant dense<5.000000e+01> : tensor<2x1xf32>
%cst_7 = stablehlo.constant dense<1.000000e-03> : tensor<1xf32>
%cst_8 = stablehlo.constant dense<3.600000e-02> : tensor<1xf32>
%cst_9 = stablehlo.constant dense<1.200000e-01> : tensor<1xf32>
%cst_10 = stablehlo.constant dense<1.250000e-01> : tensor<1xf32>
%cst_11 = stablehlo.constant dense<1.250000e-02> : tensor<1xf32>
%cst_12 = stablehlo.constant dense<0.00999999977> : tensor<1xf32>
%cst_13 = stablehlo.constant dense<5.500000e+01> : tensor<1xf32>
%cst_14 = stablehlo.constant dense<3.500000e+01> : tensor<1xf32>
%cst_15 = stablehlo.constant dense<7.000000e-02> : tensor<1xf32>
%cst_16 = stablehlo.constant dense<5.000000e-02> : tensor<1xf32>
%cst_17 = stablehlo.constant dense<-2.500000e-02> : tensor<1xf32>
%cst_18 = stablehlo.constant dense<4.000000e+00> : tensor<1xf32>
%cst_19 = stablehlo.constant dense<0.055555556> : tensor<1xf32>
%cst_20 = stablehlo.constant dense<6.500000e+01> : tensor<1xf32>
%cst_21 = stablehlo.constant dense<1.000000e+00> : tensor<1xf32>
%cst_22 = stablehlo.constant dense<2.000000e+01> : tensor<1xf32>
%cst_23 = stablehlo.constant dense<1.000000e-01> : tensor<1xf32>
%cst_24 = stablehlo.constant dense<4.000000e+01> : tensor<1xf32>
%c = stablehlo.constant dense<0> : tensor<1x1xi32>
%c_25 = stablehlo.constant dense<0> : tensor<2x1xi32>
%cst_26 = stablehlo.constant dense<1591.54944> : tensor<2xf32>
%cst_27 = stablehlo.constant dense<0.000000e+00> : tensor<1xf32>
%c_28 = stablehlo.constant dense<401> : tensor
%c_29 = stablehlo.constant dense<1> : tensor
%c_30 = stablehlo.constant dense<0> : tensor
%cst_31 = stablehlo.constant dense<1.000000e+00> : tensor
%cst_32 = stablehlo.constant dense<"0x0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000CDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3DCDCCCC3D000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"> : tensor<401x2xf32>
%cst_33 = stablehlo.constant dense<-7.000000e+01> : tensor
%cst_34 = stablehlo.constant dense<-7.000000e+01> : tensor<1xf32>
%cst_35 = stablehlo.constant dense<2.000000e-01> : tensor<1xf32>
%cst_36 = stablehlo.constant dense<0.000000e+00> : tensor<401x1xf32>
%0:6 = stablehlo.while(%iterArg = %c_30, %iterArg_37 = %cst_35, %iterArg_38 = %cst_35, %iterArg_39 = %cst_35, %iterArg_40 = %cst_34, %iterArg_41 = %cst_36) : tensor, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<401x1xf32>
cond {
%3 = stablehlo.compare LT, %iterArg, %c_28, SIGNED : (tensor, tensor) -> tensor
stablehlo.return %3 : tensor
} do {
%3 = stablehlo.dynamic_slice %cst_32, %iterArg, %c_30, sizes = [1, 2] : (tensor<401x2xf32>, tensor, tensor) -> tensor<1x2xf32>
%4 = stablehlo.reshape %3 : (tensor<1x2xf32>) -> tensor<2xf32>
%5 = stablehlo.multiply %4, %cst_26 : tensor<2xf32>
%6 = "stablehlo.scatter"(%cst_27, %c_25, %5) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter, unique_indices = false}> ({
^bb0(%arg0: tensor, %arg1: tensor):
%134 = stablehlo.add %arg0, %arg1 : tensor
stablehlo.return %134 : tensor
}) : (tensor<1xf32>, tensor<2x1xi32>, tensor<2xf32>) -> tensor<1xf32>
%7 = "stablehlo.gather"(%iterArg_38, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%8 = "stablehlo.gather"(%iterArg_37, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%9 = "stablehlo.gather"(%iterArg_39, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%10 = "stablehlo.gather"(%iterArg_40, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%11 = stablehlo.add %10, %cst_24 : tensor<1xf32>
%12 = stablehlo.negate %11 : tensor<1xf32>
%13 = stablehlo.multiply %12, %cst_23 : tensor<1xf32>
%14 = stablehlo.minimum %cst_22, %13 : tensor<1xf32>
%15 = stablehlo.exponential %14 : tensor<1xf32>
%16 = stablehlo.subtract %15, %cst_21 : tensor<1xf32>
%17 = stablehlo.divide %12, %16 : tensor<1xf32>
%18 = stablehlo.multiply %cst_23, %17 : tensor<1xf32>
%19 = stablehlo.add %10, %cst_20 : tensor<1xf32>
%20 = stablehlo.negate %19 : tensor<1xf32>
%21 = stablehlo.multiply %20, %cst_19 : tensor<1xf32>
%22 = stablehlo.minimum %cst_22, %21 : tensor<1xf32>
%23 = stablehlo.exponential %22 : tensor<1xf32>
%24 = stablehlo.multiply %cst_18, %23 : tensor<1xf32>
%25 = stablehlo.add %18, %24 : tensor<1xf32>
%26 = stablehlo.divide %cst_21, %25 : tensor<1xf32>
%27 = stablehlo.multiply %18, %26 : tensor<1xf32>
%28 = stablehlo.divide %cst_17, %26 : tensor<1xf32>
%29 = stablehlo.minimum %cst_22, %28 : tensor<1xf32>
%30 = stablehlo.exponential %29 : tensor<1xf32>
%31 = stablehlo.multiply %7, %30 : tensor<1xf32>
%32 = stablehlo.subtract %cst_21, %30 : tensor<1xf32>
%33 = stablehlo.multiply %27, %32 : tensor<1xf32>
%34 = stablehlo.add %31, %33 : tensor<1xf32>
%35 = stablehlo.multiply %20, %cst_16 : tensor<1xf32>
%36 = stablehlo.minimum %cst_22, %35 : tensor<1xf32>
%37 = stablehlo.exponential %36 : tensor<1xf32>
%38 = stablehlo.multiply %cst_15, %37 : tensor<1xf32>
%39 = stablehlo.add %10, %cst_14 : tensor<1xf32>
%40 = stablehlo.negate %39 : tensor<1xf32>
%41 = stablehlo.multiply %40, %cst_23 : tensor<1xf32>
%42 = stablehlo.minimum %cst_22, %41 : tensor<1xf32>
%43 = stablehlo.exponential %42 : tensor<1xf32>
%44 = stablehlo.add %43, %cst_21 : tensor<1xf32>
%45 = stablehlo.divide %cst_21, %44 : tensor<1xf32>
%46 = stablehlo.add %38, %45 : tensor<1xf32>
%47 = stablehlo.divide %cst_21, %46 : tensor<1xf32>
%48 = stablehlo.multiply %38, %47 : tensor<1xf32>
%49 = stablehlo.divide %cst_17, %47 : tensor<1xf32>
%50 = stablehlo.minimum %cst_22, %49 : tensor<1xf32>
%51 = stablehlo.exponential %50 : tensor<1xf32>
%52 = stablehlo.multiply %8, %51 : tensor<1xf32>
%53 = stablehlo.subtract %cst_21, %51 : tensor<1xf32>
%54 = stablehlo.multiply %48, %53 : tensor<1xf32>
%55 = stablehlo.add %52, %54 : tensor<1xf32>
%56 = stablehlo.add %10, %cst_13 : tensor<1xf32>
%57 = stablehlo.negate %56 : tensor<1xf32>
%58 = stablehlo.multiply %57, %cst_23 : tensor<1xf32>
%59 = stablehlo.minimum %cst_22, %58 : tensor<1xf32>
%60 = stablehlo.exponential %59 : tensor<1xf32>
%61 = stablehlo.subtract %60, %cst_21 : tensor<1xf32>
%62 = stablehlo.divide %57, %61 : tensor<1xf32>
%63 = stablehlo.multiply %cst_12, %62 : tensor<1xf32>
%64 = stablehlo.multiply %20, %cst_11 : tensor<1xf32>
%65 = stablehlo.minimum %cst_22, %64 : tensor<1xf32>
%66 = stablehlo.exponential %65 : tensor<1xf32>
%67 = stablehlo.multiply %cst_10, %66 : tensor<1xf32>
%68 = stablehlo.add %63, %67 : tensor<1xf32>
%69 = stablehlo.divide %cst_21, %68 : tensor<1xf32>
%70 = stablehlo.multiply %63, %69 : tensor<1xf32>
%71 = stablehlo.divide %cst_17, %69 : tensor<1xf32>
%72 = stablehlo.minimum %cst_22, %71 : tensor<1xf32>
%73 = stablehlo.exponential %72 : tensor<1xf32>
%74 = stablehlo.multiply %9, %73 : tensor<1xf32>
%75 = stablehlo.subtract %cst_21, %73 : tensor<1xf32>
%76 = stablehlo.multiply %70, %75 : tensor<1xf32>
%77 = stablehlo.add %74, %76 : tensor<1xf32>
%78 = "stablehlo.scatter"(%iterArg_38, %c, %34) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter, unique_indices = true}> ({
^bb0(%arg0: tensor, %arg1: tensor):
stablehlo.return %arg1 : tensor
}) : (tensor<1xf32>, tensor<1x1xi32>, tensor<1xf32>) -> tensor<1xf32>
%79 = "stablehlo.scatter"(%iterArg_37, %c, %55) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter, unique_indices = true}> ({
^bb0(%arg0: tensor, %arg1: tensor):
stablehlo.return %arg1 : tensor
}) : (tensor<1xf32>, tensor<1x1xi32>, tensor<1xf32>) -> tensor<1xf32>
%80 = "stablehlo.scatter"(%iterArg_39, %c, %77) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter, unique_indices = true}> ({
^bb0(%arg0: tensor, %arg1: tensor):
stablehlo.return %arg1 : tensor
}) : (tensor<1xf32>, tensor<1x1xi32>, tensor<1xf32>) -> tensor<1xf32>
%81 = "stablehlo.gather"(%78, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%82 = "stablehlo.gather"(%79, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%83 = "stablehlo.gather"(%80, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%84 = "stablehlo.gather"(%iterArg_40, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%85 = "stablehlo.gather"(%iterArg_40, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%86 = stablehlo.add %85, %cst_7 : tensor<1xf32>
%87 = stablehlo.concatenate %84, %86, dim = 0 : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>
%88 = stablehlo.reshape %87 : (tensor<2xf32>) -> tensor<2x1xf32>
%89 = stablehlo.multiply %81, %81 : tensor<1xf32>
%90 = stablehlo.multiply %89, %81 : tensor<1xf32>
%91 = stablehlo.multiply %cst_9, %90 : tensor<1xf32>
%92 = stablehlo.multiply %91, %82 : tensor<1xf32>
%93 = stablehlo.multiply %83, %83 : tensor<1xf32>
%94 = stablehlo.multiply %93, %93 : tensor<1xf32>
%95 = stablehlo.multiply %cst_8, %94 : tensor<1xf32>
%96 = stablehlo.subtract %88, %cst_6 : tensor<2x1xf32>
%97 = stablehlo.broadcast_in_dim %92, dims = [1] : (tensor<1xf32>) -> tensor<2x1xf32>
%98 = stablehlo.multiply %97, %96 : tensor<2x1xf32>
%99 = stablehlo.subtract %88, %cst_5 : tensor<2x1xf32>
%100 = stablehlo.broadcast_in_dim %95, dims = [1] : (tensor<1xf32>) -> tensor<2x1xf32>
%101 = stablehlo.multiply %100, %99 : tensor<2x1xf32>
%102 = stablehlo.add %98, %101 : tensor<2x1xf32>
%103 = stablehlo.subtract %88, %cst_4 : tensor<2x1xf32>
%104 = stablehlo.multiply %cst_3, %103 : tensor<2x1xf32>
%105 = stablehlo.add %102, %104 : tensor<2x1xf32>
%106 = stablehlo.slice %105 [1:2, 0:1] : (tensor<2x1xf32>) -> tensor<1x1xf32>
%107 = stablehlo.reshape %106 : (tensor<1x1xf32>) -> tensor<1xf32>
%108 = stablehlo.slice %105 [0:1, 0:1] : (tensor<2x1xf32>) -> tensor<1x1xf32>
%109 = stablehlo.reshape %108 : (tensor<1x1xf32>) -> tensor<1xf32>
%110 = stablehlo.subtract %107, %109 : tensor<1xf32>
%111 = stablehlo.multiply %110, %cst_2 : tensor<1xf32>
%112 = "stablehlo.gather"(%iterArg_40, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%113 = stablehlo.multiply %111, %112 : tensor<1xf32>
%114 = stablehlo.subtract %109, %113 : tensor<1xf32>
%115 = stablehlo.negate %114 : tensor<1xf32>
%116 = stablehlo.multiply %110, %cst_0 : tensor<1xf32>
%117 = stablehlo.multiply %115, %cst : tensor<1xf32>
%118 = "stablehlo.gather"(%6, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%119 = stablehlo.add %117, %118 : tensor<1xf32>
%120 = "stablehlo.scatter"(%6, %c, %119) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter, unique_indices = true}> ({
^bb0(%arg0: tensor, %arg1: tensor):
stablehlo.return %arg1 : tensor
}) : (tensor<1xf32>, tensor<1x1xi32>, tensor<1xf32>) -> tensor<1xf32>
%121 = stablehlo.multiply %cst_1, %116 : tensor<1xf32>
%122 = "stablehlo.scatter"(%cst_27, %c, %121) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter, unique_indices = true}> ({
^bb0(%arg0: tensor, %arg1: tensor):
stablehlo.return %arg1 : tensor
}) : (tensor<1xf32>, tensor<1x1xi32>, tensor<1xf32>) -> tensor<1xf32>
%123 = "stablehlo.scatter"(%122, %c, %cst_21) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter, unique_indices = true}> ({
^bb0(%arg0: tensor, %arg1: tensor):
%134 = stablehlo.add %arg0, %cst_31 : tensor
stablehlo.return %134 : tensor
}) : (tensor<1xf32>, tensor<1x1xi32>, tensor<1xf32>) -> tensor<1xf32>
%124 = "stablehlo.gather"(%iterArg_40, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%125 = "stablehlo.gather"(%120, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%126 = stablehlo.multiply %cst_1, %125 : tensor<1xf32>
%127 = stablehlo.add %124, %126 : tensor<1xf32>
%128 = "stablehlo.scatter"(%cst_27, %c, %127) <{indices_are_sorted = false, scatter_dimension_numbers = #stablehlo.scatter, unique_indices = true}> ({
^bb0(%arg0: tensor, %arg1: tensor):
stablehlo.return %arg1 : tensor
}) : (tensor<1xf32>, tensor<1x1xi32>, tensor<1xf32>) -> tensor<1xf32>
%129 = stablehlo.divide %128, %123 : tensor<1xf32>
%130 = "stablehlo.gather"(%129, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1xf32>, tensor<1x1xi32>) -> tensor<1xf32>
%131 = stablehlo.reshape %130 : (tensor<1xf32>) -> tensor<1x1xf32>
%132 = stablehlo.dynamic_update_slice %iterArg_41, %131, %iterArg, %c_30 : (tensor<401x1xf32>, tensor<1x1xf32>, tensor, tensor) -> tensor<401x1xf32>
%133 = stablehlo.add %iterArg, %c_29 {enzymexla.bounds = [[1 : i32, 401 : i32]]} : tensor
stablehlo.return %133, %79, %78, %80, %130, %132 : tensor, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>, tensor<401x1xf32>
}
%1 = stablehlo.reshape %0#5 : (tensor<401x1xf32>) -> tensor<1x401xf32>
%2 = stablehlo.pad %1, %cst_33, low = [0, 1], high = [0, 0], interior = [0, 0] : (tensor<1x401xf32>, tensor) -> tensor<1x402xf32>
return %2 : tensor<1x402xf32>
}
}
```

from jaxley https://jaxley.readthedocs.io/en/latest/index.html