func.func @main(%arg0: tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>> {
      %0 = "tf.Identity"(%arg0) : (tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
      func.return %0 : tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
    }
  })";
  SetPassThatChangedIdentity("");

      %3 = "mhlo.dynamic_iota"(%2) {iota_dimension = 0 : i64} : (tensor<1xi64>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [4]>>
      %4 = mhlo.multiply %3, %3 : tensor<?xi32, #mhlo.type_extensions<bounds = [4]>>
      return %4 : tensor<?xi32, #mhlo.type_extensions<bounds = [4]>>
    }
  })";
  CellReader<int64_t> legal_error(kFailedLegalizationStreamz);

    func.func @main(%arg0: tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>> {
      %0 = "tf.Identity"(%arg0) : (tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
      func.return %0 : tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
    }
  })";

  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
  func.func @main(%arg0: tensor<?xi32, #mhlo.type_extensions<bounds = [3]>> ) -> (tensor<?xi32, #mhlo.type_extensions<bounds = [3]>>) {
    func.return %arg0 : tensor<?xi32, #mhlo.type_extensions<bounds = [3]>>
  }
})";

  auto compilation_result = CompileMlirModule(
      kMlirModuleWithBoundedDynamicArgStr,

      return %1: tensor<2048xi32>
   }
   // CHECK-LABEL: func @tpu_func
   // CHECK: mhlo.type_extensions
   func.func @tpu_func (

    // CHECK-SAME: {dimensions = dense<0> : tensor<1xi64>}
    // CHECK-SAME: (tensor<?xf32, #mhlo.type_extensions<bounds = [4]>>) -> tensor<?xf32, #mhlo.type_extensions<bounds = [4]>>
    %1 = "tf.ReverseV2"(%0, %axis) : (tensor<?xf32>, tensor<1xi32>) -> tensor<?xf32>

    // CHECK: %[[RESULT:.*]] = tensor.cast %[[REVERSED]] : tensor<?xf32, #mhlo.type_extensions<bounds = [4]>> to tensor<?xf32>
    // CHECK: return %[[RESULT]] : tensor<?xf32>

  // CHECK-NEXT: %21 = "mhlo.set_dimension_size"(%18, %20) <{dimension = 1 : i64}> : (tensor<1x4xi32>, tensor<i32>) -> tensor<1x?xi32, #mhlo.type_extensions<bounds = [?, 4]>>
  // CHECK-NEXT: %cast = tensor.cast %21 : tensor<1x?xi32, #mhlo.type_extensions<bounds = [?, 4]>> to tensor<1x4xi32>
  // CHECK-NEXT: return %cast : tensor<1x4xi32>
  %sizes = "tf.Const"() {value = dense<[1, -1]> : tensor<2xi32>} : () -> (tensor<2xi32>)

    return %3: tensor<512xi32>
  }
func.func private @_func(%arg0: tensor<?xi32, #mhlo.type_extensions<bounds = [512]>> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) -> (tensor<512xi32>) {
    %0 = "tf.A"(%arg0) {} : (tensor<?xi32, #mhlo.type_extensions<bounds = [512]>>) -> tensor<512xi32>
    return %0 : tensor<512xi32>
  }
}

// -----

Paul Merlin <******@****.***> 1690890212 +0200