// 2. Replace the body of "tfl_wrapped_jax_random_uniform" and
//    "tfl_wrapped_jax_random_normal" with tfl.CustomOp("RandomUniform") and
//     tfl.CustomOp("RandomStandardNormal"), respectively.

#include <string>

#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Debug.h"

// CHECK:           %[[VAL_0:.*]] = stablehlo.constant dense<[3, 4]> : tensor<2xi32>
// CHECK:           %[[VAL_1:.*]] = "tfl.custom"(%[[VAL_0]]) <{custom_code = "RandomStandardNormal", custom_option = #tfl<const_bytes : "0x">}> : (tensor<2xi32>) -> tensor<3x4xf32>
// CHECK:           %[[VAL_2:.*]] = stablehlo.tuple %[[VAL_1]] : tuple<tensor<3x4xf32>>
// CHECK:           return %[[VAL_2]] : tuple<tensor<3x4xf32>>

# CHECK-SAME: tf_saved_model.initializer_type = "init_op"
# CHECK-SAME: }
# CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{{{.*dense<.*> : tensor<2xi32>.*}}}> {{{.*}}} : () -> tensor<2xi32>
# CHECK: %[[RAND_STD_NORMAL:.*]] = "tf.RandomStandardNormal"(%[[CST_0]])
# CHECK: "tf.AssignVariableOp"(%[[ARG_0]], %[[RAND_STD_NORMAL]]){{.*}}: (tensor<!tf_type.resource<tensor<1x3xf32>>>, tensor<1x3xf32>) -> ()
# CHECK: return

# The function for the signature "key".

  }

  bool IsStatefulRandomOp(absl::string_view op_name) const {
    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; };

    ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<?xf32>, %arg4: tensor<1xi32>):
      %8 = "tf.Identity"(%arg4) {device = ""} : (tensor<1xi32>) -> tensor<1xi32>
      // CHECK: tf.RandomStandardNormal{{.*}}(%[[SHAPE]])
      %9 = "tf.RandomStandardNormal"(%arg4) {device = "", seed = 87654321 : i64, seed2 = 0 : i64} : (tensor<1xi32>) -> tensor<?xf32>
      %10 = "tf.Pow"(%9, %5) {device = ""} : (tensor<?xf32>, tensor<f32>) -> tensor<?xf32>

      "QuantizeAndDequantizeV3",
      "QuantizeAndDequantizeV4",
      "RFFT",
      "RFFT2D",
      "RFFT3D",
      "RGBToHSV",
      "RandomShuffle",
      "RandomStandardNormal",
      "RandomUniform",
      "RandomUniformInt",
      "ReadVariableOp",
      "ReadVariableXlaSplitND",
      "ResizeBilinear",
      "ResizeBilinearGrad",
      "ResizeNearestNeighbor",

// CHECK:  "tfl.random_uniform"(%arg0) <{seed = 0 : i64, seed2 = 0 : i64}> : (tensor<3xi32>) -> tensor<?x?x?xf32>
}

func.func @random_standard_normal_f32(%arg0: tensor<3xi32>) -> tensor<?x?x?xf32> {
  %0 = "tf.RandomStandardNormal"(%arg0) {seed = 0 : i64, seed2 = 0: i64} : (tensor<3xi32>) -> tensor<?x?x?xf32>
  func.return %0 : tensor<?x?x?xf32>

// CHECK-LABEL:random_standard_normal_f32

  // CHECK: %[[CONV:.*]] = mhlo.convert %arg0 : (tensor<3xi32>) -> tensor<3xi64>
  // CHECK: %[[F32:.*]] = "mhlo.rng"(%[[ZERO]], %[[ONE]], %[[CONV]]) {{.*NORMAL.*}} -> tensor<12x?x64xf32>
  %0 = "tf.RandomStandardNormal"(%arg0) : (tensor<3xi32>) -> tensor<12x?x64xf32>
  // CHECK: return %[[F32]]
  func.return %0 : tensor<12x?x64xf32>
}

//===----------------------------------------------------------------------===//

    static bool isCompatibleReturnTypes(TypeRange inferred, TypeRange actual) {
      return ArraysAreCastCompatible(inferred, actual);
    }
  }];
}

def TF_RandomStandardNormalOp : TF_Op<"RandomStandardNormal", [TF_CannotDuplicate, TF_RandomGeneratorSideEffect]> {
  let summary = "Outputs random values from a normal distribution.";

  let description = [{