return op_name == "RandomUniform" || op_name == "RandomShuffle" ||
           op_name == "RandomUniformInt" || op_name == "RandomStandardNormal" ||
           op_name == "TruncatedNormal" || op_name == "Multinomial";
  }

  bool OpProducesOrConsumesVariant(const Node& node) const {
    auto is_variant = [](DataType dtype) { return dtype == DT_VARIANT; };

    // return %[[RESULT]]
    func.return %0 : tensor<*xf32>
  }

  func.func @truncated_normal() -> tensor<2x2xf32> {
    // CHECK-NOT: tf.TruncatedNormal
    %0 = mhlo.constant dense<[2, 2]> : tensor<2xi32>
    %1 = "tf.TruncatedNormal"(%0) {T = i32, device = "", dtype = f32, seed = 0 : i64, seed2 = 1950157571 : i64} : (tensor<2xi32>) -> tensor<2x2xf32>
    func.return %1 : tensor<2x2xf32>
  }