and sets the shape of these inputs to be dynamic shaped. This will ensure
    that the generated HLO program is correctly reflecting the dynamic shape.
  }];

  // Required for mhlo bounded shape extension.
  let dependentDialects = ["mhlo::MhloDialect"];
  let constructor = "TFTPU::CreateTPUAnnotateDynamicShapeInputsPass()";
}

  if (auto input_ty = mlir::dyn_cast<RankedTensorType>(old_arg_type)) {
    ArrayRef<int64_t> bounds = hlo::encodingToBounds(input_ty.getEncoding());
    // The input type has bounded dynamic dimension.
    if (!bounds.empty()) {
      SmallVector<int64_t> new_bounds(bounds.begin(), bounds.end());
      SmallVector<int64_t> new_shape(shape.begin(), shape.end());