[XLA] Enhance the HLO verifier to report errors for illegal Rng instructions.

Modify ShapeVerifier::HandleRng to detect illegal Rng instructions.

Add test cases to test the handling of Rng instructions in the HLO verifier.

Modify the document for XLA random operations to reflect the supported data
types.

PiperOrigin-RevId: 208220062

1 files changed, 7 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_verifier.h b/tensorflow/compiler/xla/service/hlo_verifier.h
index 5a56a44f35..c942fab08e 100644
--- a/tensorflow/compiler/xla/service/hlo_verifier.h
+++ b/tensorflow/compiler/xla/service/hlo_verifier.h
@@ -106,6 +106,13 @@ class ShapeVerifier : public DfsHloVisitor {
   Status CheckVariadicShape(const HloInstruction* instruction);
 
  private:
+  // Return true if the shapes of the two operands have the same element type,
+  // and the result shape either has the same element type as the operand
+  // shapes or mixed precision is allowed and the result shape and the operand
+  // shapes have floating point element types.
+  bool HasCompatibleElementTypes(const Shape& shape_0, const Shape& shape_1,
+                                 const Shape& result_shape);
+
   // Whether the inputs and output of an instruction can contain both F32s and
   // BF16s. Tuples that include both F32s and BF16s are allowed regardless of
   // this flag.