See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// This file implements logic for legalizing HLO to TensorFlow.

#include <cassert>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <functional>
#include <memory>
#include <numeric>
#include <optional>
#include <string>

  Type new_arg_type = tensorflow::GetTypeFromTFTensorShape(shape, element_type);

  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());

  let description = [{
    This pass looks for the usage of the result of TPUCopyWithDynamicShapeOp
    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"];