RankedTensorType index_type = tensorflow::GetTypeFromTFTensorShape(
        input_type.getShape(), index_element_type);

    std::optional<int64_t> optional_axis =
        GetIntegerHLOAxisFromTFAxis(op.getDimension(), input_type.getRank());
    if (!optional_axis.has_value())
      return rewriter.notifyMatchFailure(op, "required axis");
    int64_t axis = optional_axis.value();

    IntegerAttr iota_dimension =