| 372 | } |
| 373 | |
| 374 | Status ShapeVerifier::HandleRng(HloInstruction* instruction) { |
| 375 | TF_RETURN_IF_ERROR(CheckOperandCount(instruction, 2)); |
| 376 | |
| 377 | const Shape& shape_0 = instruction->operand(0)->shape(); |
| 378 | const Shape& shape_1 = instruction->operand(1)->shape(); |
| 379 | if (!ShapeUtil::IsScalar(shape_0) || !ShapeUtil::IsScalar(shape_1)) { |
| 380 | return InternalError( |
| 381 | "Expected scalar types for the two operands of Rng instruction: %s", |
| 382 | instruction->ToString()); |
| 383 | } |
| 384 | |
| 385 | if (!HasCompatibleElementTypes(shape_0, shape_1, instruction->shape())) { |
| 386 | return InternalError( |
| 387 | "Expected compatible element types for the result and the two operands" |
| 388 | " of Rng instruction: %s", |
| 389 | instruction->ToString()); |
| 390 | } |
| 391 | |
| 392 | PrimitiveType element_type = shape_0.element_type(); |
| 393 | switch (instruction->random_distribution()) { |
| 394 | case RNG_UNIFORM: |
| 395 | if (!primitive_util::IsFloatingPointType(element_type) && |
| 396 | !primitive_util::IsIntegralType(element_type) && |
| 397 | element_type != PRED) { |
| 398 | return InternalError( |
| 399 | "Element type not supported." |
| 400 | " Expected element to be of floating point type, integral type or" |
| 401 | " predicate type for RngUniform: %s", |
| 402 | instruction->ToString()); |
| 403 | } |
| 404 | break; |
| 405 | |
| 406 | case RNG_NORMAL: |
| 407 | if (!primitive_util::IsFloatingPointType(element_type)) { |
| 408 | return InternalError( |
| 409 | "Element type not supported." |
| 410 | " Expected element to be FloatingPointType for RngNormal: %s", |
| 411 | instruction->ToString()); |
| 412 | } |
| 413 | break; |
| 414 | default: |
| 415 | return InternalError( |
| 416 | "Invalid Rng distribution %s", |
| 417 | RandomDistribution_Name(instruction->random_distribution())); |
| 418 | } |
| 419 | |
| 420 | return Status::OK(); |
| 421 | } |
| 422 | |
| 423 | Status ShapeVerifier::HandleRngBitGenerator(HloInstruction* hlo) { |
| 424 | if (!hlo->shape().IsTuple() || hlo->shape().tuple_shapes_size() != 2) { |
no test coverage detected