}

    Value mul_rhs_value;
    if (!output_type.hasRank() || (output_type.getNumDynamicDims() > 0)) {
      auto status_or_const_op =
          CreateConstOpWithSingleValue(&rewriter, loc, input_type, 1);
      if (!status_or_const_op.ok()) {
        return failure();
      }

      mul_rhs_value = rewriter.create<TFL::FillOp>(