// 1. IotaOp.
// 2. IotaOp + BroadCastInDim.
// 3. IotaOp + Reshape.
// 4. Constant (folded Iota) + BroadCastInDim.
// 5. Constant (folded result).
// Moreover, the dimensions has to match the iota_dimension.
bool MatchIota(DenseIntElementsAttr dimensions, Value iota);
}  // namespace odml
}  // namespace mlir

    if (iota_init.getValues<APInt>()[0] != 0) return failure();

    // Verify that the second argument is an Iota op along the same dimension
    // as the reduction.
    Value iota = reduce_op.getInputs().back();
    if (!MatchIota(reduce_op.getDimensions(), iota)) return failure();

    // Match the reduction computation.
    const bool is_float = mlir::isa<FloatType>(operand_init.getElementType());