}

  if (!IsOfRankOrUnranked(op.getDepth(), 0)) {
    return op.emitOpError() << "requires depth to be a scalar";
  }
  if (!IsOfRankOrUnranked(op.getOnValue(), 0)) {
    return op.emitOpError() << "requires on_value to be a scalar";
  }
  if (!IsOfRankOrUnranked(op.getOffValue(), 0)) {
    return op.emitOpError() << "requires off_value to be a scalar";
  }

    Value compare = rewriter.create<mhlo::CompareOp>(
        loc, broadcast_indices, iota, ComparisonDirection::EQ);
    Value on_value = rewriter.create<BroadcastOp>(
        loc, op.getType(), op.getOnValue(),
        GetI64ElementsAttr(output_dims, &rewriter));
    Value off_value = rewriter.create<BroadcastOp>(
        loc, op.getType(), op.getOffValue(),
        GetI64ElementsAttr(output_dims, &rewriter));