func.func @not_eliminate_multinomial(%0: tensor<*xf32>, %1: tensor<*xi32>) -> (tensor<*xi64>, tensor<*xi64>) {
  // CHECK-NEXT: tf.Multinomial
  // CHECK-SAME: seed = 0
  // CHECK-NEXT: tf.Multinomial
  // CHECK-SAME: seed = 1
  // CHECK-NEXT: tf.Multinomial
  // CHECK-SAME: seed = 0
  // CHECK-NEXT: return

    func.return %1 : tensor<1000xi32>
  }

  // CHECK-LABEL: multinomial
  func.func @multinomial(%arg0: tensor<2x4xf32>, %seed: tensor<i32>, %seed2: tensor<i32>) -> tensor<2x10xi32> {
    // CHECK-NOT: tf.Multinomial
    %samples = "tf.Const"() { value = dense<10> : tensor<i32> } : () -> tensor<i32>
    %1 = "tf.Multinomial"(%arg0, %samples) {seed = 0, seed2 = 0}: (tensor<2x4xf32>, tensor<i32>) -> tensor<2x10xi32>

    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 absl::c_any_of(node.input_types(), is_variant) ||

}

func.func @multinomial_i64(%arg0: tensor<2xf32>, %arg1: tensor<1xi32>) -> tensor<10xi64> {
  %0 = "tf.Multinomial"(%arg0, %arg1) {seed = 0 : i64, seed2 = 0: i64} : (tensor<2xf32>, tensor<1xi32>) -> tensor<10xi64>
  func.return %0 : tensor<10xi64>

// CHECK-LABEL:multinomial_i64
// CHECK: "tfl.multinomial"(%arg0, %arg1) <{seed = 0 : i64, seed2 = 0 : i64}> : (tensor<2xf32>, tensor<1xi32>) -> tensor<10xi64>
}

  CONV_3D_TRANSPOSE = 141,
  VAR_HANDLE = 142,
  READ_VARIABLE = 143,
  ASSIGN_VARIABLE = 144,
  BROADCAST_ARGS = 145,
  RANDOM_STANDARD_NORMAL = 146,
  BUCKETIZE = 147,
  RANDOM_UNIFORM = 148,
  MULTINOMIAL = 149,
  GELU = 150,
  DYNAMIC_UPDATE_SLICE = 151,
  RELU_0_TO_1 = 152,
  UNSORTED_SEGMENT_PROD = 153,
  UNSORTED_SEGMENT_MAX = 154,
  UNSORTED_SEGMENT_SUM = 155,
  ATAN2 = 156,

}

// -----

// CHECK-LABEL: testMultinomial
func.func @testMultinomial(%arg0: tensor<2xf32>, %arg1: tensor<1xi32>) -> tensor<10xi64> {
  // CHECK: "tfl.multinomial"(%arg0, %arg1)
  %0 = "tfl.multinomial"(%arg0, %arg1) {seed = 0 : i64, seed2 = 0: i64} : (tensor<2xf32>, tensor<1xi32>) -> tensor<10xi64>
  func.return %0 :  tensor<10xi64>
}

// -----

// CHECK-LABEL: testMultinomialInt32

      "MaxPool3D",
      "MaxPool3DGrad",
      "MaxPool3DGradGrad",
      "MaxPoolGrad",
      "MaxPoolGradGrad",
      "MaxPoolGradGradV2",
      "MaxPoolGradV2",
      "MaxPoolV2",
      "Multinomial",
      "NextAfter",
      "NonMaxSuppressionV3",
      "NonMaxSuppressionV4",
      "ParallelDynamicStitch",
      "ParameterizedTruncatedNormal",
      "PartitionedCall",
      "Polygamma",

  );

  let results = (outs
    TFL_TensorOf<[F32]>:$out);
  let hasOptions = 1;
  let customOption = "RandomOptions";
}

def TFL_MultinomialOp : TFL_Op<"multinomial", []> {
  let summary = "Draws samples from a categorical distribution.";

  let description = [{
The generated values will have a categorical distribution based on the `logits`

  );

  let results = (outs
    Res<TF_StrTensor, [{A string representing the resource.}]>:$string_handle
  );
}

def TF_MultinomialOp : TF_Op<"Multinomial", [TF_CannotDuplicate]> {
  let summary = "Draws samples from a multinomial distribution.";

  let arguments = (ins
    Arg<TF_IntOrFpTensor, [{2-D Tensor with shape `[batch_size, num_classes]`.  Each slice `[i, :]`

        `TransformedDistribution`.
    *   Fix a bug in `import_meta_graph`'s handling of partitioned variables
        when
    *   Ensure `tf.distributions.Multinomial` doesn't underflow in `log_prob`.
        Before this change, all partitions of an integer variable were
        initialized with the shape of the unpartitioned variable; after this