| 1454 | } // namespace |
| 1455 | |
| 1456 | bool LiteralBase::IsAll(int8 value) const { |
| 1457 | return root_piece().ForEachSubpieceWithBool([&](const ShapeIndex& index, |
| 1458 | const Piece& piece) { |
| 1459 | if (!piece.subshape().IsArray()) { |
| 1460 | return true; |
| 1461 | } |
| 1462 | |
| 1463 | auto piece_is_all = [&]() { |
| 1464 | switch (shape().element_type()) { |
| 1465 | case U8: |
| 1466 | if (value >= 0) { |
| 1467 | return AllElementsEqualValue<uint8>(piece.data<uint8>(), value); |
| 1468 | } |
| 1469 | return false; |
| 1470 | case U16: |
| 1471 | if (value >= 0) { |
| 1472 | return AllElementsEqualValue<uint16>(piece.data<uint16>(), value); |
| 1473 | } |
| 1474 | return false; |
| 1475 | case U32: |
| 1476 | if (value >= 0) { |
| 1477 | return AllElementsEqualValue<uint32>(piece.data<uint32>(), value); |
| 1478 | } |
| 1479 | return false; |
| 1480 | case U64: |
| 1481 | if (value >= 0) { |
| 1482 | return AllElementsEqualValue<uint64>(piece.data<uint64>(), value); |
| 1483 | } |
| 1484 | return false; |
| 1485 | case S8: |
| 1486 | return AllElementsEqualValue<int8>(piece.data<int8>(), value); |
| 1487 | case S16: |
| 1488 | return AllElementsEqualValue<int16>(piece.data<int16>(), value); |
| 1489 | case S32: |
| 1490 | return AllElementsEqualValue<int32>(piece.data<int32>(), value); |
| 1491 | case S64: |
| 1492 | return AllElementsEqualValue<int64>(piece.data<int64>(), value); |
| 1493 | case F32: |
| 1494 | return AllElementsEqualValue<float>(piece.data<float>(), value); |
| 1495 | case F64: |
| 1496 | return AllElementsEqualValue<double>(piece.data<double>(), value); |
| 1497 | case F16: |
| 1498 | return AllElementsEqualValue<half>(piece.data<half>(), |
| 1499 | static_cast<half>(value)); |
| 1500 | case BF16: |
| 1501 | return AllElementsEqualValue<bfloat16>(piece.data<bfloat16>(), |
| 1502 | static_cast<bfloat16>(value)); |
| 1503 | case PRED: |
| 1504 | if (value == 0) { |
| 1505 | return AllElementsEqualValue<bool>(piece.data<bool>(), false); |
| 1506 | } |
| 1507 | if (value == 1) { |
| 1508 | return AllElementsEqualValue<bool>(piece.data<bool>(), true); |
| 1509 | } |
| 1510 | return false; |
| 1511 | default: |
| 1512 | return false; |
| 1513 | } |