}

// TODO(lyandy): Investigate supported dtypes (`minval`, `maxval`, `output`) for
// `tf.StatefulUniformInt`. tf2xla kernels support i32, i64, ui32, and ui64
// while TensorFlow CPU/GPU kernels only support i32 and i64.
def TF_StatefulUniformIntOp : TF_Op<"StatefulUniformInt", []> {
  let summary = "Outputs random integers from a uniform distribution.";

  let description = [{

#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

#define DEBUG_TYPE "tf-tfl-legalization"

namespace mlir {
namespace TFL {
namespace {
// Returns a TF_CastOp to I32. This function is used for CastOps that are
// intermediate nodes in a TableGen pattern result. In such a case, the
// destination type is not inferred and must be given explicitly.
//

// Same as above, but if indices are already 1D, then no reshape is needed.
def ReplaceOneHotDynamicShapeFullyConnectedWithLookup : Pat<
  (TFL_FullyConnectedOp:$outputs
    (TFL_OneHotOp
      1DTensorOf<[I32]>:$indices,
      (Arith_ConstantOp $depth),
      (Arith_ConstantOp ConstantAttr<RankedF32ElementsAttr<[]>, "1.0f">),
      (Arith_ConstantOp ConstantAttr<RankedF32ElementsAttr<[]>, "0.0f">),