Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

stablehlo-to-tensorrt does not support converting stablehlo.dynamic_gather #264

Open
parthchadha opened this issue Oct 10, 2024 · 0 comments
Open

stablehlo-to-tensorrt does not support converting stablehlo.dynamic_gather #264

parthchadha opened this issue Oct 10, 2024 · 0 comments
Labels
mlir-tensorrt Pull request for the mlir-tensorrt project

Comments

@parthchadha
Copy link
Collaborator 

Loaded TensorRT version 10.4.0.26 but compiled for TensorRT 10.2.0.19. This can result in crashes or unintended behavior.
    failed to load TensorRT builder timing cache from provided file path (/tmp/tripy-cache), error on opening file: cannot open input file '/tmp/tripy-cache': No such file or directory, the compiler will start with a fresh empty cache
    (t475)error: op: %48 = "stablehlo.dynamic_gather"(%47#0, %47#2, %47#3) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>}> : (tensor<?x?xf32>, tensor<?xi32>, tensor<2xi32>) -> tensor<?x?xf32> from function main is invalid, post clustering.
    (t475)error: op: %49 = "stablehlo.dynamic_iota"(%47#1) <{iota_dimension = 1 : i64}> : (tensor<2xi32>) -> tensor<?x?xi32> from function main is invalid, post clustering.
    (t475)error: op: %61 = "stablehlo.dynamic_gather"(%60#0, %60#1, %60#2) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1], index_vector_dim = 1>}> : (tensor<?x?xf32>, tensor<?xi32>, tensor<2xi32>) -> tensor<?x?xf32> from function main is invalid, post clustering.

    This error occured while trying to compile the following FlatIR expression:
          |
          | t475: [rank=(2), shape=((-1, -1)), dtype=(float32), loc=(gpu:0)] = DynamicGatherOp(t429, t474, t_inter637, axis=0)
          | 


    Note: This originated from the following expression:

    --> /tripy/tripy/frontend/trace/ops/slice.py:239 in __getitem__()
          |
      239 |                 result = gather(result, dim, idx)
          |                          ^^^^^^^^^^^^^^^^^^^^^^^^

module @ins_all_mask_logits_all_iou_scores_outs_t614_4 {
  func.func @main(%arg0: tensor<?x?x?x?xf32> {tensorrt.shape_profile = #tensorrt.shape_profile<min = [2, 1, 4, 5], opt = [2, 2, 4, 5], max = [2, 3, 4, 5]>}, %arg1: tensor<?x?xf32> {tensorrt.shape_profile = #tensorrt.shape_profile<min = [1, 3], opt = [2, 3], max = [3, 3]>}) -> tensor<?x?x?xf32> {
    %0 = stablehlo.get_dimension_size %arg1, dim = 0 : (tensor<?x?xf32>) -> tensor<i32>
    %c = stablehlo.constant dense<1> : tensor<1xi32>
    %1 = stablehlo.reshape %0 : (tensor<i32>) -> tensor<1xi32>
    %2 = stablehlo.get_dimension_size %arg1, dim = 1 : (tensor<?x?xf32>) -> tensor<i32>
    %3 = stablehlo.reshape %2 : (tensor<i32>) -> tensor<1xi32>
    %4 = stablehlo.concatenate %1, %3, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %c_0 = stablehlo.constant dense<0> : tensor<i32>
    %c_1 = stablehlo.constant dense<1> : tensor<i32>
    %c_2 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_3 = stablehlo.constant dense<1> : tensor<1xi32>
    %5 = stablehlo.compare  LE, %c_2, %c_3 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %6 = stablehlo.select %5, %c_2, %c_3 : tensor<1xi1>, tensor<1xi32>
    %c_4 = stablehlo.constant dense<1> : tensor<1xi32>
    %7 = stablehlo.real_dynamic_slice %4, %6, %c_3, %c_4 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %c_5 = stablehlo.constant dense<> : tensor<0xi32>
    %8 = stablehlo.dynamic_reshape %7, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %c_6 = stablehlo.constant dense<> : tensor<0xi32>
    %9 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %10 = stablehlo.select %9, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %11 = stablehlo.dynamic_broadcast_in_dim %8, %10, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %12 = stablehlo.dynamic_broadcast_in_dim %c_0, %10, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %13 = stablehlo.compare  LT, %11, %12 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_7 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_8 = stablehlo.constant dense<1> : tensor<1xi32>
    %14 = stablehlo.compare  LE, %c_7, %c_8 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %15 = stablehlo.select %14, %c_7, %c_8 : tensor<1xi1>, tensor<1xi32>
    %c_9 = stablehlo.constant dense<1> : tensor<1xi32>
    %16 = stablehlo.real_dynamic_slice %4, %15, %c_8, %c_9 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %17 = stablehlo.dynamic_reshape %16, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %c_10 = stablehlo.constant dense<0> : tensor<1xi32>
    %18 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %19 = stablehlo.dynamic_broadcast_in_dim %13, %18, dims = [] : (tensor<i1>, tensor<1xi32>) -> tensor<?xi1>
    %20 = stablehlo.get_dimension_size %19, dim = 0 : (tensor<?xi1>) -> tensor<i32>
    %21 = stablehlo.reshape %20 : (tensor<i32>) -> tensor<1xi32>
    %22 = stablehlo.compare  EQ, %21, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %23 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %24 = stablehlo.dynamic_broadcast_in_dim %17, %23, dims = [] : (tensor<i32>, tensor<1xi32>) -> tensor<?xi32>
    %25 = stablehlo.get_dimension_size %24, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %26 = stablehlo.reshape %25 : (tensor<i32>) -> tensor<1xi32>
    %27 = stablehlo.select %22, %26, %21 : tensor<1xi1>, tensor<1xi32>
    %28 = stablehlo.compare  EQ, %27, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %29 = stablehlo.get_dimension_size %24, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %30 = stablehlo.reshape %29 : (tensor<i32>) -> tensor<1xi32>
    %31 = stablehlo.compare  EQ, %30, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_11 = stablehlo.constant dense<1> : tensor<i32>
    %c_12 = stablehlo.constant dense<1> : tensor<1xi32>
    %32 = stablehlo.select %31, %c_12, %30 : tensor<1xi1>, tensor<1xi32>
    %33 = stablehlo.select %28, %32, %27 : tensor<1xi1>, tensor<1xi32>
    %34 = stablehlo.dynamic_broadcast_in_dim %19, %33, dims = [0] : (tensor<?xi1>, tensor<1xi32>) -> tensor<?xi1>
    %35 = stablehlo.dynamic_broadcast_in_dim %24, %33, dims = [0] : (tensor<?xi32>, tensor<1xi32>) -> tensor<?xi32>
    %36 = stablehlo.dynamic_broadcast_in_dim %c_10, %33, dims = [0] : (tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %37 = stablehlo.select %34, %35, %36 : tensor<?xi1>, tensor<?xi32>
    %c_13 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_14 = stablehlo.constant dense<1> : tensor<1xi32>
    %38 = stablehlo.compare  LE, %c_13, %c_14 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %39 = stablehlo.select %38, %c_13, %c_14 : tensor<1xi1>, tensor<1xi32>
    %c_15 = stablehlo.constant dense<1> : tensor<1xi32>
    %40 = stablehlo.real_dynamic_slice %4, %39, %c_14, %c_15 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %41 = stablehlo.dynamic_reshape %40, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %42 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %43 = stablehlo.select %42, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %44 = stablehlo.dynamic_broadcast_in_dim %41, %43, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %45 = stablehlo.dynamic_broadcast_in_dim %c_0, %43, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %46 = stablehlo.compare  GE, %44, %45 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_16 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_17 = stablehlo.constant dense<1> : tensor<1xi32>
    %47 = stablehlo.compare  LE, %c_16, %c_17 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %48 = stablehlo.select %47, %c_16, %c_17 : tensor<1xi1>, tensor<1xi32>
    %c_18 = stablehlo.constant dense<1> : tensor<1xi32>
    %49 = stablehlo.real_dynamic_slice %4, %48, %c_17, %c_18 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %50 = stablehlo.dynamic_reshape %49, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %51 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %52 = stablehlo.select %51, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %53 = stablehlo.dynamic_broadcast_in_dim %50, %52, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %54 = stablehlo.dynamic_broadcast_in_dim %41, %52, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %55 = stablehlo.add %53, %54 : tensor<i32>
    %56 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %57 = stablehlo.select %56, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %58 = stablehlo.compare  EQ, %57, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %59 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %60 = stablehlo.select %59, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %61 = stablehlo.select %58, %60, %57 : tensor<0xi1>, tensor<0xi32>
    %62 = stablehlo.dynamic_broadcast_in_dim %46, %61, dims = [] : (tensor<i1>, tensor<0xi32>) -> tensor<i1>
    %63 = stablehlo.dynamic_broadcast_in_dim %41, %61, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %64 = stablehlo.dynamic_broadcast_in_dim %55, %61, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %65 = stablehlo.select %62, %63, %64 : tensor<i1>, tensor<i32>
    %66 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %67 = stablehlo.select %66, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %68 = stablehlo.dynamic_broadcast_in_dim %65, %67, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %69 = stablehlo.dynamic_broadcast_in_dim %c_0, %67, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %70 = stablehlo.compare  LT, %68, %69 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_19 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_20 = stablehlo.constant dense<1> : tensor<1xi32>
    %71 = stablehlo.compare  LE, %c_19, %c_20 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %72 = stablehlo.select %71, %c_19, %c_20 : tensor<1xi1>, tensor<1xi32>
    %c_21 = stablehlo.constant dense<1> : tensor<1xi32>
    %73 = stablehlo.real_dynamic_slice %4, %72, %c_20, %c_21 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %74 = stablehlo.dynamic_reshape %73, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %75 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %76 = stablehlo.select %75, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %77 = stablehlo.dynamic_broadcast_in_dim %74, %76, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %78 = stablehlo.dynamic_broadcast_in_dim %65, %76, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %79 = stablehlo.compare  LT, %77, %78 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_22 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_23 = stablehlo.constant dense<1> : tensor<1xi32>
    %80 = stablehlo.compare  LE, %c_22, %c_23 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %81 = stablehlo.select %80, %c_22, %c_23 : tensor<1xi1>, tensor<1xi32>
    %c_24 = stablehlo.constant dense<1> : tensor<1xi32>
    %82 = stablehlo.real_dynamic_slice %4, %81, %c_23, %c_24 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %83 = stablehlo.dynamic_reshape %82, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %84 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %85 = stablehlo.select %84, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %86 = stablehlo.compare  EQ, %85, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %87 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %88 = stablehlo.select %87, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %89 = stablehlo.select %86, %88, %85 : tensor<0xi1>, tensor<0xi32>
    %90 = stablehlo.dynamic_broadcast_in_dim %79, %89, dims = [] : (tensor<i1>, tensor<0xi32>) -> tensor<i1>
    %91 = stablehlo.dynamic_broadcast_in_dim %83, %89, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %92 = stablehlo.dynamic_broadcast_in_dim %65, %89, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %93 = stablehlo.select %90, %91, %92 : tensor<i1>, tensor<i32>
    %94 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %95 = stablehlo.dynamic_broadcast_in_dim %70, %94, dims = [] : (tensor<i1>, tensor<1xi32>) -> tensor<?xi1>
    %96 = stablehlo.get_dimension_size %95, dim = 0 : (tensor<?xi1>) -> tensor<i32>
    %97 = stablehlo.reshape %96 : (tensor<i32>) -> tensor<1xi32>
    %98 = stablehlo.compare  EQ, %97, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_25 = stablehlo.constant dense<1> : tensor<i32>
    %c_26 = stablehlo.constant dense<1> : tensor<1xi32>
    %99 = stablehlo.select %98, %c_26, %97 : tensor<1xi1>, tensor<1xi32>
    %100 = stablehlo.compare  EQ, %99, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_27 = stablehlo.constant dense<1> : tensor<i32>
    %c_28 = stablehlo.constant dense<1> : tensor<1xi32>
    %101 = stablehlo.compare  EQ, %c_28, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %102 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %103 = stablehlo.dynamic_broadcast_in_dim %93, %102, dims = [] : (tensor<i32>, tensor<1xi32>) -> tensor<?xi32>
    %104 = stablehlo.get_dimension_size %103, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %105 = stablehlo.reshape %104 : (tensor<i32>) -> tensor<1xi32>
    %106 = stablehlo.select %101, %105, %c_28 : tensor<1xi1>, tensor<1xi32>
    %107 = stablehlo.select %100, %106, %99 : tensor<1xi1>, tensor<1xi32>
    %108 = stablehlo.dynamic_broadcast_in_dim %95, %107, dims = [0] : (tensor<?xi1>, tensor<1xi32>) -> tensor<?xi1>
    %109 = stablehlo.dynamic_broadcast_in_dim %c_10, %107, dims = [0] : (tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %110 = stablehlo.dynamic_broadcast_in_dim %103, %107, dims = [0] : (tensor<?xi32>, tensor<1xi32>) -> tensor<?xi32>
    %111 = stablehlo.select %108, %109, %110 : tensor<?xi1>, tensor<?xi32>
    %c_29 = stablehlo.constant dense<2> : tensor<i32>
    %c_30 = stablehlo.constant dense<1> : tensor<1xi32>
    %c_31 = stablehlo.constant dense<2> : tensor<1xi32>
    %112 = stablehlo.compare  LE, %c_30, %c_31 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %113 = stablehlo.select %112, %c_30, %c_31 : tensor<1xi1>, tensor<1xi32>
    %c_32 = stablehlo.constant dense<1> : tensor<1xi32>
    %114 = stablehlo.real_dynamic_slice %4, %113, %c_31, %c_32 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %115 = stablehlo.dynamic_reshape %114, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %116 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %117 = stablehlo.select %116, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %118 = stablehlo.dynamic_broadcast_in_dim %115, %117, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %119 = stablehlo.dynamic_broadcast_in_dim %c_1, %117, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %120 = stablehlo.compare  LT, %118, %119 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_33 = stablehlo.constant dense<1> : tensor<1xi32>
    %c_34 = stablehlo.constant dense<2> : tensor<1xi32>
    %121 = stablehlo.compare  LE, %c_33, %c_34 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %122 = stablehlo.select %121, %c_33, %c_34 : tensor<1xi1>, tensor<1xi32>
    %c_35 = stablehlo.constant dense<1> : tensor<1xi32>
    %123 = stablehlo.real_dynamic_slice %4, %122, %c_34, %c_35 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %124 = stablehlo.dynamic_reshape %123, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %125 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %126 = stablehlo.dynamic_broadcast_in_dim %120, %125, dims = [] : (tensor<i1>, tensor<1xi32>) -> tensor<?xi1>
    %127 = stablehlo.get_dimension_size %126, dim = 0 : (tensor<?xi1>) -> tensor<i32>
    %128 = stablehlo.reshape %127 : (tensor<i32>) -> tensor<1xi32>
    %129 = stablehlo.compare  EQ, %128, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %130 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %131 = stablehlo.dynamic_broadcast_in_dim %124, %130, dims = [] : (tensor<i32>, tensor<1xi32>) -> tensor<?xi32>
    %132 = stablehlo.get_dimension_size %131, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %133 = stablehlo.reshape %132 : (tensor<i32>) -> tensor<1xi32>
    %134 = stablehlo.select %129, %133, %128 : tensor<1xi1>, tensor<1xi32>
    %135 = stablehlo.compare  EQ, %134, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %136 = stablehlo.get_dimension_size %131, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %137 = stablehlo.reshape %136 : (tensor<i32>) -> tensor<1xi32>
    %138 = stablehlo.compare  EQ, %137, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_36 = stablehlo.constant dense<1> : tensor<i32>
    %c_37 = stablehlo.constant dense<1> : tensor<1xi32>
    %139 = stablehlo.select %138, %c_37, %137 : tensor<1xi1>, tensor<1xi32>
    %140 = stablehlo.select %135, %139, %134 : tensor<1xi1>, tensor<1xi32>
    %141 = stablehlo.dynamic_broadcast_in_dim %126, %140, dims = [0] : (tensor<?xi1>, tensor<1xi32>) -> tensor<?xi1>
    %142 = stablehlo.dynamic_broadcast_in_dim %131, %140, dims = [0] : (tensor<?xi32>, tensor<1xi32>) -> tensor<?xi32>
    %143 = stablehlo.dynamic_broadcast_in_dim %c, %140, dims = [0] : (tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %144 = stablehlo.select %141, %142, %143 : tensor<?xi1>, tensor<?xi32>
    %c_38 = stablehlo.constant dense<1> : tensor<1xi32>
    %c_39 = stablehlo.constant dense<2> : tensor<1xi32>
    %145 = stablehlo.compare  LE, %c_38, %c_39 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %146 = stablehlo.select %145, %c_38, %c_39 : tensor<1xi1>, tensor<1xi32>
    %c_40 = stablehlo.constant dense<1> : tensor<1xi32>
    %147 = stablehlo.real_dynamic_slice %4, %146, %c_39, %c_40 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %148 = stablehlo.dynamic_reshape %147, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %149 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %150 = stablehlo.select %149, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %151 = stablehlo.dynamic_broadcast_in_dim %148, %150, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %152 = stablehlo.dynamic_broadcast_in_dim %c_0, %150, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %153 = stablehlo.compare  GE, %151, %152 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_41 = stablehlo.constant dense<1> : tensor<1xi32>
    %c_42 = stablehlo.constant dense<2> : tensor<1xi32>
    %154 = stablehlo.compare  LE, %c_41, %c_42 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %155 = stablehlo.select %154, %c_41, %c_42 : tensor<1xi1>, tensor<1xi32>
    %c_43 = stablehlo.constant dense<1> : tensor<1xi32>
    %156 = stablehlo.real_dynamic_slice %4, %155, %c_42, %c_43 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %157 = stablehlo.dynamic_reshape %156, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %158 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %159 = stablehlo.select %158, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %160 = stablehlo.dynamic_broadcast_in_dim %157, %159, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %161 = stablehlo.dynamic_broadcast_in_dim %148, %159, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %162 = stablehlo.add %160, %161 : tensor<i32>
    %163 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %164 = stablehlo.select %163, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %165 = stablehlo.compare  EQ, %164, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %166 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %167 = stablehlo.select %166, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %168 = stablehlo.select %165, %167, %164 : tensor<0xi1>, tensor<0xi32>
    %169 = stablehlo.dynamic_broadcast_in_dim %153, %168, dims = [] : (tensor<i1>, tensor<0xi32>) -> tensor<i1>
    %170 = stablehlo.dynamic_broadcast_in_dim %148, %168, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %171 = stablehlo.dynamic_broadcast_in_dim %162, %168, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %172 = stablehlo.select %169, %170, %171 : tensor<i1>, tensor<i32>
    %173 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %174 = stablehlo.select %173, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %175 = stablehlo.dynamic_broadcast_in_dim %172, %174, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %176 = stablehlo.dynamic_broadcast_in_dim %c_0, %174, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %177 = stablehlo.compare  LT, %175, %176 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_44 = stablehlo.constant dense<1> : tensor<1xi32>
    %c_45 = stablehlo.constant dense<2> : tensor<1xi32>
    %178 = stablehlo.compare  LE, %c_44, %c_45 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %179 = stablehlo.select %178, %c_44, %c_45 : tensor<1xi1>, tensor<1xi32>
    %c_46 = stablehlo.constant dense<1> : tensor<1xi32>
    %180 = stablehlo.real_dynamic_slice %4, %179, %c_45, %c_46 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %181 = stablehlo.dynamic_reshape %180, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %182 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %183 = stablehlo.select %182, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %184 = stablehlo.dynamic_broadcast_in_dim %181, %183, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %185 = stablehlo.dynamic_broadcast_in_dim %172, %183, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %186 = stablehlo.compare  LT, %184, %185 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_47 = stablehlo.constant dense<1> : tensor<1xi32>
    %c_48 = stablehlo.constant dense<2> : tensor<1xi32>
    %187 = stablehlo.compare  LE, %c_47, %c_48 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %188 = stablehlo.select %187, %c_47, %c_48 : tensor<1xi1>, tensor<1xi32>
    %c_49 = stablehlo.constant dense<1> : tensor<1xi32>
    %189 = stablehlo.real_dynamic_slice %4, %188, %c_48, %c_49 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %190 = stablehlo.dynamic_reshape %189, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %191 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %192 = stablehlo.select %191, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %193 = stablehlo.compare  EQ, %192, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %194 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %195 = stablehlo.select %194, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %196 = stablehlo.select %193, %195, %192 : tensor<0xi1>, tensor<0xi32>
    %197 = stablehlo.dynamic_broadcast_in_dim %186, %196, dims = [] : (tensor<i1>, tensor<0xi32>) -> tensor<i1>
    %198 = stablehlo.dynamic_broadcast_in_dim %190, %196, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %199 = stablehlo.dynamic_broadcast_in_dim %172, %196, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %200 = stablehlo.select %197, %198, %199 : tensor<i1>, tensor<i32>
    %201 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %202 = stablehlo.dynamic_broadcast_in_dim %177, %201, dims = [] : (tensor<i1>, tensor<1xi32>) -> tensor<?xi1>
    %203 = stablehlo.get_dimension_size %202, dim = 0 : (tensor<?xi1>) -> tensor<i32>
    %204 = stablehlo.reshape %203 : (tensor<i32>) -> tensor<1xi32>
    %205 = stablehlo.compare  EQ, %204, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_50 = stablehlo.constant dense<1> : tensor<i32>
    %c_51 = stablehlo.constant dense<1> : tensor<1xi32>
    %206 = stablehlo.select %205, %c_51, %204 : tensor<1xi1>, tensor<1xi32>
    %207 = stablehlo.compare  EQ, %206, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_52 = stablehlo.constant dense<1> : tensor<i32>
    %c_53 = stablehlo.constant dense<1> : tensor<1xi32>
    %208 = stablehlo.compare  EQ, %c_53, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %209 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %210 = stablehlo.dynamic_broadcast_in_dim %200, %209, dims = [] : (tensor<i32>, tensor<1xi32>) -> tensor<?xi32>
    %211 = stablehlo.get_dimension_size %210, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %212 = stablehlo.reshape %211 : (tensor<i32>) -> tensor<1xi32>
    %213 = stablehlo.select %208, %212, %c_53 : tensor<1xi1>, tensor<1xi32>
    %214 = stablehlo.select %207, %213, %206 : tensor<1xi1>, tensor<1xi32>
    %215 = stablehlo.dynamic_broadcast_in_dim %202, %214, dims = [0] : (tensor<?xi1>, tensor<1xi32>) -> tensor<?xi1>
    %216 = stablehlo.dynamic_broadcast_in_dim %c_10, %214, dims = [0] : (tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %217 = stablehlo.dynamic_broadcast_in_dim %210, %214, dims = [0] : (tensor<?xi32>, tensor<1xi32>) -> tensor<?xi32>
    %218 = stablehlo.select %215, %216, %217 : tensor<?xi1>, tensor<?xi32>
    %219 = stablehlo.reshape %37 : (tensor<?xi32>) -> tensor<1xi32>
    %220 = stablehlo.reshape %111 : (tensor<?xi32>) -> tensor<1xi32>
    %221 = stablehlo.compare  LE, %219, %220 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %222 = stablehlo.select %221, %219, %220 : tensor<1xi1>, tensor<1xi32>
    %223 = stablehlo.reshape %144 : (tensor<?xi32>) -> tensor<1xi32>
    %224 = stablehlo.reshape %218 : (tensor<?xi32>) -> tensor<1xi32>
    %225 = stablehlo.compare  LE, %223, %224 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %226 = stablehlo.select %225, %223, %224 : tensor<1xi1>, tensor<1xi32>
    %227 = stablehlo.concatenate %222, %226, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %228 = stablehlo.concatenate %220, %224, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %c_54 = stablehlo.constant dense<1> : tensor<1xi32>
    %c_55 = stablehlo.constant dense<1> : tensor<1xi32>
    %229 = stablehlo.concatenate %c_54, %c_55, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %230 = stablehlo.real_dynamic_slice %arg1, %227, %228, %229 : (tensor<?x?xf32>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<?x?xf32>
    %231 = stablehlo.get_dimension_size %230, dim = 0 : (tensor<?x?xf32>) -> tensor<i32>
    %232 = stablehlo.reshape %231 : (tensor<i32>) -> tensor<1xi32>
    %233 = stablehlo.get_dimension_size %230, dim = 1 : (tensor<?x?xf32>) -> tensor<i32>
    %234 = stablehlo.reshape %233 : (tensor<i32>) -> tensor<1xi32>
    %235 = stablehlo.concatenate %232, %234, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %c_56 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_57 = stablehlo.constant dense<1> : tensor<1xi32>
    %236 = stablehlo.compare  LE, %c_56, %c_57 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %237 = stablehlo.select %236, %c_56, %c_57 : tensor<1xi1>, tensor<1xi32>
    %c_58 = stablehlo.constant dense<1> : tensor<1xi32>
    %238 = stablehlo.real_dynamic_slice %235, %237, %c_57, %c_58 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %239 = stablehlo.dynamic_reshape %238, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %240 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %241 = stablehlo.select %240, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %242 = stablehlo.dynamic_broadcast_in_dim %c_0, %241, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %243 = stablehlo.dynamic_broadcast_in_dim %239, %241, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %244 = stablehlo.subtract %242, %243 : tensor<i32>
    %245 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %246 = stablehlo.select %245, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %247 = stablehlo.dynamic_broadcast_in_dim %244, %246, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %248 = stablehlo.dynamic_broadcast_in_dim %c_1, %246, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %249 = stablehlo.divide %247, %248 : tensor<i32>
    %250 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %251 = stablehlo.select %250, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %252 = stablehlo.dynamic_broadcast_in_dim %c_0, %251, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %253 = stablehlo.dynamic_broadcast_in_dim %249, %251, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %254 = stablehlo.subtract %252, %253 : tensor<i32>
    %255 = stablehlo.dynamic_broadcast_in_dim %254, %c, dims = [] : (tensor<i32>, tensor<1xi32>) -> tensor<1xi32>
    %256 = stablehlo.dynamic_iota %255, dim = 0 : (tensor<1xi32>) -> tensor<?xf32>
    %257 = stablehlo.dynamic_broadcast_in_dim %254, %c, dims = [] : (tensor<i32>, tensor<1xi32>) -> tensor<1xi32>
    %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
    %258 = stablehlo.dynamic_broadcast_in_dim %cst, %257, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
    %259 = stablehlo.get_dimension_size %256, dim = 0 : (tensor<?xf32>) -> tensor<i32>
    %260 = stablehlo.reshape %259 : (tensor<i32>) -> tensor<1xi32>
    %261 = stablehlo.compare  EQ, %260, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %262 = stablehlo.get_dimension_size %258, dim = 0 : (tensor<?xf32>) -> tensor<i32>
    %263 = stablehlo.reshape %262 : (tensor<i32>) -> tensor<1xi32>
    %264 = stablehlo.select %261, %263, %260 : tensor<1xi1>, tensor<1xi32>
    %265 = stablehlo.dynamic_broadcast_in_dim %256, %264, dims = [0] : (tensor<?xf32>, tensor<1xi32>) -> tensor<?xf32>
    %266 = stablehlo.dynamic_broadcast_in_dim %258, %264, dims = [0] : (tensor<?xf32>, tensor<1xi32>) -> tensor<?xf32>
    %267 = stablehlo.multiply %265, %266 : tensor<?xf32>
    %268 = stablehlo.dynamic_broadcast_in_dim %254, %c, dims = [] : (tensor<i32>, tensor<1xi32>) -> tensor<1xi32>
    %cst_59 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %269 = stablehlo.dynamic_broadcast_in_dim %cst_59, %268, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
    %270 = stablehlo.get_dimension_size %267, dim = 0 : (tensor<?xf32>) -> tensor<i32>
    %271 = stablehlo.reshape %270 : (tensor<i32>) -> tensor<1xi32>
    %272 = stablehlo.compare  EQ, %271, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %273 = stablehlo.get_dimension_size %269, dim = 0 : (tensor<?xf32>) -> tensor<i32>
    %274 = stablehlo.reshape %273 : (tensor<i32>) -> tensor<1xi32>
    %275 = stablehlo.select %272, %274, %271 : tensor<1xi1>, tensor<1xi32>
    %276 = stablehlo.dynamic_broadcast_in_dim %267, %275, dims = [0] : (tensor<?xf32>, tensor<1xi32>) -> tensor<?xf32>
    %277 = stablehlo.dynamic_broadcast_in_dim %269, %275, dims = [0] : (tensor<?xf32>, tensor<1xi32>) -> tensor<?xf32>
    %278 = stablehlo.add %276, %277 : tensor<?xf32>
    %279 = stablehlo.convert %278 : (tensor<?xf32>) -> tensor<?xi32>
    %280 = stablehlo.get_dimension_size %230, dim = 0 : (tensor<?x?xf32>) -> tensor<i32>
    %281 = stablehlo.reshape %280 : (tensor<i32>) -> tensor<1xi32>
    %282 = stablehlo.get_dimension_size %230, dim = 1 : (tensor<?x?xf32>) -> tensor<i32>
    %283 = stablehlo.reshape %282 : (tensor<i32>) -> tensor<1xi32>
    %284 = stablehlo.concatenate %281, %283, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %c_60 = stablehlo.constant dense<2> : tensor<1xi32>
    %285 = stablehlo.real_dynamic_slice %284, %c, %c_60, %c : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %286 = stablehlo.concatenate %c, %285, dim = 0 : (tensor<1xi32>, tensor<?xi32>) -> tensor<2xi32>
    %287 = "stablehlo.dynamic_gather"(%230, %279, %286) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>}> : (tensor<?x?xf32>, tensor<?xi32>, tensor<2xi32>) -> tensor<?x?xf32>
    %288 = stablehlo.get_dimension_size %230, dim = 0 : (tensor<?x?xf32>) -> tensor<i32>
    %289 = stablehlo.reshape %288 : (tensor<i32>) -> tensor<1xi32>
    %290 = stablehlo.get_dimension_size %230, dim = 1 : (tensor<?x?xf32>) -> tensor<i32>
    %291 = stablehlo.reshape %290 : (tensor<i32>) -> tensor<1xi32>
    %292 = stablehlo.concatenate %289, %291, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %293 = stablehlo.dynamic_iota %292, dim = 1 : (tensor<2xi32>) -> tensor<?x?xi32>
    %294:2 = stablehlo.reduce(%230 init: %cst_59), (%293 init: %c_0) across dimensions = [1] : (tensor<?x?xf32>, tensor<?x?xi32>, tensor<f32>, tensor<i32>) -> (tensor<?xf32>, tensor<?xi32>)
     reducer(%arg2: tensor<f32>, %arg4: tensor<f32>) (%arg3: tensor<i32>, %arg5: tensor<i32>)  {
      %493 = stablehlo.compare  GE, %arg2, %arg4 : (tensor<f32>, tensor<f32>) -> tensor<i1>
      %494 = stablehlo.select %493, %arg2, %arg4 : tensor<i1>, tensor<f32>
      %495 = stablehlo.compare  EQ, %arg2, %arg4 : (tensor<f32>, tensor<f32>) -> tensor<i1>
      %496 = stablehlo.minimum %arg3, %arg5 : tensor<i32>
      %497 = stablehlo.select %493, %arg3, %arg5 : tensor<i1>, tensor<i32>
      %498 = stablehlo.select %495, %496, %497 : tensor<i1>, tensor<i32>
      stablehlo.return %494, %498 : tensor<f32>, tensor<i32>
    }
    %295 = stablehlo.get_dimension_size %287, dim = 0 : (tensor<?x?xf32>) -> tensor<i32>
    %296 = stablehlo.reshape %295 : (tensor<i32>) -> tensor<1xi32>
    %297 = stablehlo.get_dimension_size %287, dim = 1 : (tensor<?x?xf32>) -> tensor<i32>
    %298 = stablehlo.reshape %297 : (tensor<i32>) -> tensor<1xi32>
    %299 = stablehlo.concatenate %296, %298, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %300 = stablehlo.real_dynamic_slice %299, %c_10, %c, %c : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
    %301 = stablehlo.concatenate %300, %c, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %302 = "stablehlo.dynamic_gather"(%287, %294#1, %301) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1], index_vector_dim = 1>}> : (tensor<?x?xf32>, tensor<?xi32>, tensor<2xi32>) -> tensor<?x?xf32>
    %303 = stablehlo.get_dimension_size %302, dim = 0 : (tensor<?x?xf32>) -> tensor<i32>
    %304 = stablehlo.reshape %303 : (tensor<i32>) -> tensor<1xi32>
    %305 = stablehlo.get_dimension_size %302, dim = 1 : (tensor<?x?xf32>) -> tensor<i32>
    %306 = stablehlo.reshape %305 : (tensor<i32>) -> tensor<1xi32>
    %307 = stablehlo.concatenate %304, %306, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
    %c_61 = stablehlo.constant dense<2> : tensor<i32>
    %c_62 = stablehlo.constant dense<2> : tensor<1xi32>
    %c_63 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_64 = stablehlo.constant dense<1> : tensor<1xi32>
    %308 = stablehlo.compare  LE, %c_63, %c_64 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %309 = stablehlo.select %308, %c_63, %c_64 : tensor<1xi1>, tensor<1xi32>
    %c_65 = stablehlo.constant dense<1> : tensor<1xi32>
    %310 = stablehlo.real_dynamic_slice %c_62, %309, %c_64, %c_65 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %311 = stablehlo.dynamic_reshape %310, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %312 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %313 = stablehlo.select %312, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %314 = stablehlo.dynamic_broadcast_in_dim %311, %313, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %315 = stablehlo.dynamic_broadcast_in_dim %c_0, %313, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %316 = stablehlo.compare  LT, %314, %315 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_66 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_67 = stablehlo.constant dense<1> : tensor<1xi32>
    %317 = stablehlo.compare  LE, %c_66, %c_67 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %318 = stablehlo.select %317, %c_66, %c_67 : tensor<1xi1>, tensor<1xi32>
    %c_68 = stablehlo.constant dense<1> : tensor<1xi32>
    %319 = stablehlo.real_dynamic_slice %c_62, %318, %c_67, %c_68 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %320 = stablehlo.dynamic_reshape %319, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %321 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %322 = stablehlo.dynamic_broadcast_in_dim %316, %321, dims = [] : (tensor<i1>, tensor<1xi32>) -> tensor<?xi1>
    %323 = stablehlo.get_dimension_size %322, dim = 0 : (tensor<?xi1>) -> tensor<i32>
    %324 = stablehlo.reshape %323 : (tensor<i32>) -> tensor<1xi32>
    %325 = stablehlo.compare  EQ, %324, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %326 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %327 = stablehlo.dynamic_broadcast_in_dim %320, %326, dims = [] : (tensor<i32>, tensor<1xi32>) -> tensor<?xi32>
    %328 = stablehlo.get_dimension_size %327, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %329 = stablehlo.reshape %328 : (tensor<i32>) -> tensor<1xi32>
    %330 = stablehlo.select %325, %329, %324 : tensor<1xi1>, tensor<1xi32>
    %331 = stablehlo.compare  EQ, %330, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %332 = stablehlo.get_dimension_size %327, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %333 = stablehlo.reshape %332 : (tensor<i32>) -> tensor<1xi32>
    %334 = stablehlo.compare  EQ, %333, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_69 = stablehlo.constant dense<1> : tensor<i32>
    %c_70 = stablehlo.constant dense<1> : tensor<1xi32>
    %335 = stablehlo.select %334, %c_70, %333 : tensor<1xi1>, tensor<1xi32>
    %336 = stablehlo.select %331, %335, %330 : tensor<1xi1>, tensor<1xi32>
    %337 = stablehlo.dynamic_broadcast_in_dim %322, %336, dims = [0] : (tensor<?xi1>, tensor<1xi32>) -> tensor<?xi1>
    %338 = stablehlo.dynamic_broadcast_in_dim %327, %336, dims = [0] : (tensor<?xi32>, tensor<1xi32>) -> tensor<?xi32>
    %339 = stablehlo.dynamic_broadcast_in_dim %c_10, %336, dims = [0] : (tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %340 = stablehlo.select %337, %338, %339 : tensor<?xi1>, tensor<?xi32>
    %c_71 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_72 = stablehlo.constant dense<1> : tensor<1xi32>
    %341 = stablehlo.compare  LE, %c_71, %c_72 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %342 = stablehlo.select %341, %c_71, %c_72 : tensor<1xi1>, tensor<1xi32>
    %c_73 = stablehlo.constant dense<1> : tensor<1xi32>
    %343 = stablehlo.real_dynamic_slice %c_62, %342, %c_72, %c_73 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %344 = stablehlo.dynamic_reshape %343, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %345 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %346 = stablehlo.select %345, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %347 = stablehlo.dynamic_broadcast_in_dim %344, %346, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %348 = stablehlo.dynamic_broadcast_in_dim %c_1, %346, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %349 = stablehlo.compare  LT, %347, %348 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_74 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_75 = stablehlo.constant dense<1> : tensor<1xi32>
    %350 = stablehlo.compare  LE, %c_74, %c_75 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %351 = stablehlo.select %350, %c_74, %c_75 : tensor<1xi1>, tensor<1xi32>
    %c_76 = stablehlo.constant dense<1> : tensor<1xi32>
    %352 = stablehlo.real_dynamic_slice %c_62, %351, %c_75, %c_76 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %353 = stablehlo.dynamic_reshape %352, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %354 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %355 = stablehlo.dynamic_broadcast_in_dim %349, %354, dims = [] : (tensor<i1>, tensor<1xi32>) -> tensor<?xi1>
    %356 = stablehlo.get_dimension_size %355, dim = 0 : (tensor<?xi1>) -> tensor<i32>
    %357 = stablehlo.reshape %356 : (tensor<i32>) -> tensor<1xi32>
    %358 = stablehlo.compare  EQ, %357, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %359 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %360 = stablehlo.dynamic_broadcast_in_dim %353, %359, dims = [] : (tensor<i32>, tensor<1xi32>) -> tensor<?xi32>
    %361 = stablehlo.get_dimension_size %360, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %362 = stablehlo.reshape %361 : (tensor<i32>) -> tensor<1xi32>
    %363 = stablehlo.select %358, %362, %357 : tensor<1xi1>, tensor<1xi32>
    %364 = stablehlo.compare  EQ, %363, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %365 = stablehlo.get_dimension_size %360, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %366 = stablehlo.reshape %365 : (tensor<i32>) -> tensor<1xi32>
    %367 = stablehlo.compare  EQ, %366, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_77 = stablehlo.constant dense<1> : tensor<i32>
    %c_78 = stablehlo.constant dense<1> : tensor<1xi32>
    %368 = stablehlo.select %367, %c_78, %366 : tensor<1xi1>, tensor<1xi32>
    %369 = stablehlo.select %364, %368, %363 : tensor<1xi1>, tensor<1xi32>
    %370 = stablehlo.dynamic_broadcast_in_dim %355, %369, dims = [0] : (tensor<?xi1>, tensor<1xi32>) -> tensor<?xi1>
    %371 = stablehlo.dynamic_broadcast_in_dim %360, %369, dims = [0] : (tensor<?xi32>, tensor<1xi32>) -> tensor<?xi32>
    %372 = stablehlo.dynamic_broadcast_in_dim %c, %369, dims = [0] : (tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %373 = stablehlo.select %370, %371, %372 : tensor<?xi1>, tensor<?xi32>
    %374 = stablehlo.reshape %340 : (tensor<?xi32>) -> tensor<1xi32>
    %375 = stablehlo.reshape %373 : (tensor<?xi32>) -> tensor<1xi32>
    %376 = stablehlo.compare  LE, %374, %375 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %377 = stablehlo.select %376, %374, %375 : tensor<1xi1>, tensor<1xi32>
    %c_79 = stablehlo.constant dense<1> : tensor<1xi32>
    %378 = stablehlo.real_dynamic_slice %307, %377, %375, %c_79 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
    %c_80 = stablehlo.constant dense<2> : tensor<i32>
    %c_81 = stablehlo.constant dense<2> : tensor<1xi32>
    %c_82 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_83 = stablehlo.constant dense<1> : tensor<1xi32>
    %379 = stablehlo.compare  LE, %c_82, %c_83 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %380 = stablehlo.select %379, %c_82, %c_83 : tensor<1xi1>, tensor<1xi32>
    %c_84 = stablehlo.constant dense<1> : tensor<1xi32>
    %381 = stablehlo.real_dynamic_slice %c_81, %380, %c_83, %c_84 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %382 = stablehlo.dynamic_reshape %381, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %383 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %384 = stablehlo.select %383, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %385 = stablehlo.dynamic_broadcast_in_dim %382, %384, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %386 = stablehlo.dynamic_broadcast_in_dim %c_1, %384, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %387 = stablehlo.compare  LT, %385, %386 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_85 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_86 = stablehlo.constant dense<1> : tensor<1xi32>
    %388 = stablehlo.compare  LE, %c_85, %c_86 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %389 = stablehlo.select %388, %c_85, %c_86 : tensor<1xi1>, tensor<1xi32>
    %c_87 = stablehlo.constant dense<1> : tensor<1xi32>
    %390 = stablehlo.real_dynamic_slice %c_81, %389, %c_86, %c_87 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %391 = stablehlo.dynamic_reshape %390, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %392 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %393 = stablehlo.dynamic_broadcast_in_dim %387, %392, dims = [] : (tensor<i1>, tensor<1xi32>) -> tensor<?xi1>
    %394 = stablehlo.get_dimension_size %393, dim = 0 : (tensor<?xi1>) -> tensor<i32>
    %395 = stablehlo.reshape %394 : (tensor<i32>) -> tensor<1xi32>
    %396 = stablehlo.compare  EQ, %395, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %397 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %398 = stablehlo.dynamic_broadcast_in_dim %391, %397, dims = [] : (tensor<i32>, tensor<1xi32>) -> tensor<?xi32>
    %399 = stablehlo.get_dimension_size %398, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %400 = stablehlo.reshape %399 : (tensor<i32>) -> tensor<1xi32>
    %401 = stablehlo.select %396, %400, %395 : tensor<1xi1>, tensor<1xi32>
    %402 = stablehlo.compare  EQ, %401, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %403 = stablehlo.get_dimension_size %398, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %404 = stablehlo.reshape %403 : (tensor<i32>) -> tensor<1xi32>
    %405 = stablehlo.compare  EQ, %404, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_88 = stablehlo.constant dense<1> : tensor<i32>
    %c_89 = stablehlo.constant dense<1> : tensor<1xi32>
    %406 = stablehlo.select %405, %c_89, %404 : tensor<1xi1>, tensor<1xi32>
    %407 = stablehlo.select %402, %406, %401 : tensor<1xi1>, tensor<1xi32>
    %408 = stablehlo.dynamic_broadcast_in_dim %393, %407, dims = [0] : (tensor<?xi1>, tensor<1xi32>) -> tensor<?xi1>
    %409 = stablehlo.dynamic_broadcast_in_dim %398, %407, dims = [0] : (tensor<?xi32>, tensor<1xi32>) -> tensor<?xi32>
    %410 = stablehlo.dynamic_broadcast_in_dim %c, %407, dims = [0] : (tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %411 = stablehlo.select %408, %409, %410 : tensor<?xi1>, tensor<?xi32>
    %c_90 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_91 = stablehlo.constant dense<1> : tensor<1xi32>
    %412 = stablehlo.compare  LE, %c_90, %c_91 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %413 = stablehlo.select %412, %c_90, %c_91 : tensor<1xi1>, tensor<1xi32>
    %c_92 = stablehlo.constant dense<1> : tensor<1xi32>
    %414 = stablehlo.real_dynamic_slice %c_81, %413, %c_91, %c_92 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %415 = stablehlo.dynamic_reshape %414, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %416 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %417 = stablehlo.select %416, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %418 = stablehlo.dynamic_broadcast_in_dim %415, %417, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %419 = stablehlo.dynamic_broadcast_in_dim %c_0, %417, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %420 = stablehlo.compare  GE, %418, %419 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_93 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_94 = stablehlo.constant dense<1> : tensor<1xi32>
    %421 = stablehlo.compare  LE, %c_93, %c_94 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %422 = stablehlo.select %421, %c_93, %c_94 : tensor<1xi1>, tensor<1xi32>
    %c_95 = stablehlo.constant dense<1> : tensor<1xi32>
    %423 = stablehlo.real_dynamic_slice %c_81, %422, %c_94, %c_95 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %424 = stablehlo.dynamic_reshape %423, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %425 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %426 = stablehlo.select %425, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %427 = stablehlo.dynamic_broadcast_in_dim %424, %426, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %428 = stablehlo.dynamic_broadcast_in_dim %415, %426, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %429 = stablehlo.add %427, %428 : tensor<i32>
    %430 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %431 = stablehlo.select %430, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %432 = stablehlo.compare  EQ, %431, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %433 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %434 = stablehlo.select %433, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %435 = stablehlo.select %432, %434, %431 : tensor<0xi1>, tensor<0xi32>
    %436 = stablehlo.dynamic_broadcast_in_dim %420, %435, dims = [] : (tensor<i1>, tensor<0xi32>) -> tensor<i1>
    %437 = stablehlo.dynamic_broadcast_in_dim %415, %435, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %438 = stablehlo.dynamic_broadcast_in_dim %429, %435, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %439 = stablehlo.select %436, %437, %438 : tensor<i1>, tensor<i32>
    %440 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %441 = stablehlo.select %440, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %442 = stablehlo.dynamic_broadcast_in_dim %439, %441, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %443 = stablehlo.dynamic_broadcast_in_dim %c_0, %441, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %444 = stablehlo.compare  LT, %442, %443 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_96 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_97 = stablehlo.constant dense<1> : tensor<1xi32>
    %445 = stablehlo.compare  LE, %c_96, %c_97 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %446 = stablehlo.select %445, %c_96, %c_97 : tensor<1xi1>, tensor<1xi32>
    %c_98 = stablehlo.constant dense<1> : tensor<1xi32>
    %447 = stablehlo.real_dynamic_slice %c_81, %446, %c_97, %c_98 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %448 = stablehlo.dynamic_reshape %447, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %449 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %450 = stablehlo.select %449, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %451 = stablehlo.dynamic_broadcast_in_dim %448, %450, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %452 = stablehlo.dynamic_broadcast_in_dim %439, %450, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %453 = stablehlo.compare  LT, %451, %452 : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %c_99 = stablehlo.constant dense<0> : tensor<1xi32>
    %c_100 = stablehlo.constant dense<1> : tensor<1xi32>
    %454 = stablehlo.compare  LE, %c_99, %c_100 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %455 = stablehlo.select %454, %c_99, %c_100 : tensor<1xi1>, tensor<1xi32>
    %c_101 = stablehlo.constant dense<1> : tensor<1xi32>
    %456 = stablehlo.real_dynamic_slice %c_81, %455, %c_100, %c_101 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %457 = stablehlo.dynamic_reshape %456, %c_5 : (tensor<?xi32>, tensor<0xi32>) -> tensor<i32>
    %458 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %459 = stablehlo.select %458, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %460 = stablehlo.compare  EQ, %459, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %461 = stablehlo.compare  EQ, %c_6, %c_5 : (tensor<0xi32>, tensor<0xi32>) -> tensor<0xi1>
    %462 = stablehlo.select %461, %c_6, %c_6 : tensor<0xi1>, tensor<0xi32>
    %463 = stablehlo.select %460, %462, %459 : tensor<0xi1>, tensor<0xi32>
    %464 = stablehlo.dynamic_broadcast_in_dim %453, %463, dims = [] : (tensor<i1>, tensor<0xi32>) -> tensor<i1>
    %465 = stablehlo.dynamic_broadcast_in_dim %457, %463, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %466 = stablehlo.dynamic_broadcast_in_dim %439, %463, dims = [] : (tensor<i32>, tensor<0xi32>) -> tensor<i32>
    %467 = stablehlo.select %464, %465, %466 : tensor<i1>, tensor<i32>
    %468 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %469 = stablehlo.dynamic_broadcast_in_dim %444, %468, dims = [] : (tensor<i1>, tensor<1xi32>) -> tensor<?xi1>
    %470 = stablehlo.get_dimension_size %469, dim = 0 : (tensor<?xi1>) -> tensor<i32>
    %471 = stablehlo.reshape %470 : (tensor<i32>) -> tensor<1xi32>
    %472 = stablehlo.compare  EQ, %471, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_102 = stablehlo.constant dense<1> : tensor<i32>
    %c_103 = stablehlo.constant dense<1> : tensor<1xi32>
    %473 = stablehlo.select %472, %c_103, %471 : tensor<1xi1>, tensor<1xi32>
    %474 = stablehlo.compare  EQ, %473, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %c_104 = stablehlo.constant dense<1> : tensor<i32>
    %c_105 = stablehlo.constant dense<1> : tensor<1xi32>
    %475 = stablehlo.compare  EQ, %c_105, %c : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %476 = stablehlo.concatenate %c, %c_6, dim = 0 : (tensor<1xi32>, tensor<0xi32>) -> tensor<1xi32>
    %477 = stablehlo.dynamic_broadcast_in_dim %467, %476, dims = [] : (tensor<i32>, tensor<1xi32>) -> tensor<?xi32>
    %478 = stablehlo.get_dimension_size %477, dim = 0 : (tensor<?xi32>) -> tensor<i32>
    %479 = stablehlo.reshape %478 : (tensor<i32>) -> tensor<1xi32>
    %480 = stablehlo.select %475, %479, %c_105 : tensor<1xi1>, tensor<1xi32>
    %481 = stablehlo.select %474, %480, %473 : tensor<1xi1>, tensor<1xi32>
    %482 = stablehlo.dynamic_broadcast_in_dim %469, %481, dims = [0] : (tensor<?xi1>, tensor<1xi32>) -> tensor<?xi1>
    %483 = stablehlo.dynamic_broadcast_in_dim %c_10, %481, dims = [0] : (tensor<1xi32>, tensor<1xi32>) -> tensor<?xi32>
    %484 = stablehlo.dynamic_broadcast_in_dim %477, %481, dims = [0] : (tensor<?xi32>, tensor<1xi32>) -> tensor<?xi32>
    %485 = stablehlo.select %482, %483, %484 : tensor<?xi1>, tensor<?xi32>
    %486 = stablehlo.reshape %411 : (tensor<?xi32>) -> tensor<1xi32>
    %487 = stablehlo.reshape %485 : (tensor<?xi32>) -> tensor<1xi32>
    %488 = stablehlo.compare  LE, %486, %487 : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
    %489 = stablehlo.select %488, %486, %487 : tensor<1xi1>, tensor<1xi32>
    %c_106 = stablehlo.constant dense<1> : tensor<1xi32>
    %490 = stablehlo.real_dynamic_slice %307, %489, %487, %c_106 : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
    %491 = stablehlo.concatenate %378, %c, %490, dim = 0 : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32>
    %492 = stablehlo.dynamic_broadcast_in_dim %302, %491, dims = [0, 2] : (tensor<?x?xf32>, tensor<3xi32>) -> tensor<?x?x?xf32>
    return %492 : tensor<?x?x?xf32>
  }
}
@parthchadha parthchadha added the mlir-tensorrt Pull request for the mlir-tensorrt project label Oct 10, 2024
@christopherbate christopherbate changed the title dynamic_gather op can not be lowered into TensorRT stablehlo-to-tensorrt does not support converting stablehlo.dynamic_gather Nov 5, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
mlir-tensorrt Pull request for the mlir-tensorrt project
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
1 participant
@parthchadha