%14 = stablehlo.subtract %12, %13 : tensor<8x16x16xi32>  // q1 - z1
    %15 = stablehlo.convert %11 : (tensor<8x16x4xi8>) -> tensor<8x16x4xi32>
    %16 = stablehlo.broadcast_in_dim %6, dims = [0, 1, 2] : (tensor<1x1x1xi32>) -> tensor<8x16x4xi32>
    %17 = stablehlo.subtract %15, %16 : tensor<8x16x4xi32>  // q2 - z2
    // Corresponds to einsum expression: b i j, b j d -> b i d

    %zero = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
    %elem_1 = "tf._SomeOtherOp"() : () -> tensor<8x1xf32>
    %tl_3 = "tf.TensorListSetItem"(%tl_1, %zero, %elem_1) : (tensor<!tf_type.variant<tensor<?x1xf32>>>, tensor<i32>, tensor<8x1xf32>) -> tensor<!tf_type.variant<tensor<?x1xf32>>>
    func.return
  }

  // CHECK-LABEL: multiple_mutation_dynamic_element_shape

  %1 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<16x8xf32>) -> tensor<8x16xf32>
  %2 = "stablehlo.reduce_window"(%1, %0) ({
  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
      %3 = stablehlo.maximum %arg1, %arg2 : tensor<f32>
      stablehlo.return %3 : tensor<f32>
  }) {window_dimensions = array<i64: 2, 2>, window_strides = array<i64: 2, 2>} : (tensor<8x16xf32>, tensor<f32>) -> tensor<4x8xf32>
  return %2 : tensor<4x8xf32>
}

// CHECK:           %[[VAL_5:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64>
// CHECK:           %[[VAL_6:.*]] = "tf.Transpose"(%[[VAL_4]], %[[VAL_5]]) : (tensor<1x1xi32>, tensor<2xi64>) -> tensor<1x1xi32>
// CHECK:           return %[[VAL_6]] : tensor<1x1xi32>
// CHECK:         }

    "tfl.yield"(%1#0, %1#1, %1#2) : (tensor<1x?x16xf32>, tensor<1x!tf_type.string>, tensor<1xi32>) -> ()
  }) {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x!tf_type.string>) -> (tensor<1x?x16xf32>, tensor<1x!tf_type.string>, tensor<1xi32>)

  %cst = arith.constant dense<[[1], [2]]> : tensor<2x1xi32>

  // CHECK: %[[CST:.*]] = arith.constant dense<2> : tensor<1xi32>
  // CHECK: return %[[CST]]
  %0 = "tfl.rank"(%cst) : (tensor<2x1xi32>) -> tensor<1xi32>
  func.return %0 : tensor<1xi32>
}

// CHECK-LABEL: @rank_input_known_rank
func.func @rank_input_known_rank(%arg0 : tensor<2x1xi32>) -> tensor<1xi32> {

  func.return %1 : tensor<16xi32>
  // CHECK-DAG: %[[CST:.*]] = arith.constant dense<-1> : tensor<1xi32>
  // CHECK: %[[ARG_MIN:.*]] = "tfl.arg_min"(%arg0, %[[CST]]) : (tensor<16x1024xf32>, tensor<1xi32>) -> tensor<16xi32>
  // CHECK: return %[[ARG_MIN]] : tensor<16xi32>
}

// CHECK-LABEL: DontRemoveSoftmaxNegativeBetaBeforeArgmax
func.func @DontRemoveSoftmaxNegativeBetaBeforeArgmax(%arg0: tensor<16x1024xf32>) -> tensor<16xi32> {

  // CHECK:  %5 = "tfl.arg_max"(%arg0, %cst) : (tensor<1x32x1xf32>, tensor<1xi32>) -> tensor<1x1xi32>
  // CHECK:  return %4, %5 : tensor<1x1xf32>, tensor<1x1xi32>
}

// CHECK-LABEL:   func @convert_pytorch_argmax
func.func @convert_pytorch_argmax(%arg0: tensor<1x9xi32>) -> tensor<1xi32> {
  %0 = mhlo.constant dense<0> : tensor<i32>

    // CHECK-NOT: tf.Multinomial
    %samples = "tf.Const"() { value = dense<10> : tensor<i32> } : () -> tensor<i32>
    %1 = "tf.Multinomial"(%arg0, %samples) {seed = 0, seed2 = 0}: (tensor<2x4xf32>, tensor<i32>) -> tensor<2x10xi32>
    func.return %1 : tensor<2x10xi32>
  }

  // CHECK-LABEL: @set_dynamic_dimension_size

  %cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi32>} : () -> tensor<16xi32>
  %cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>