func.func @main(%arg0: tensor<3xi64> {tf._user_specified_name = "resource", tf.aliasing_output = 3 : i64}) -> () attributes {tf.entry_function = {control_outputs = "stateful_normal/RngReadAndSkip,stateful_uniform/RngReadAndSkip,stateful_uniform_full_int/RngReadAndSkip", inputs = "stateful_normal_rngreadandskip_resource", outputs = "identity_RetVal,identity_1_RetVal,identity_2_RetVal"}} {

      .cast<TensorType>()
      .getElementType()
      .cast<ResourceType>()
      .getSubtypes()
      .front();
}

// Decompose tf.RngReadAndSkip.
//
// For Philox, the resource variable holds a tensor<3xi64> with the state:
//   [counter_lo, counter_hi, key]
//
//   RngReadAndSkip increments the 128 bit counter value by 256 * delta and
//   returns the original state value.
//

  func.func @main(%arg0: tensor<3xi64> {tf._user_specified_name = "resource", tf.aliasing_output = 3 : i64}) -> () attributes {tf.entry_function = {control_outputs = "stateful_normal/RngReadAndSkip,stateful_uniform/RngReadAndSkip,stateful_uniform_full_int/RngReadAndSkip", inputs = "stateful_normal_rngreadandskip_resource", outputs = "identity_RetVal,identity_1_RetVal,identity_2_RetVal"}} {

}


// Test that tf.RngReadAndSkip op is decomposed.
// CHECK-LABEL: func @decompose_rng_read_and_skip_op
func.func @decompose_rng_read_and_skip_op(%resource: tensor<!tf_type.resource<tensor<3xi64>>>) -> tensor<3xi64> {
  %0 = "tf_device.cluster"() ({
    // We rely on the TensorFlow StatefulRandomOpsTest to check it is lowered
    // correctly.
    // CHECK-NOT: tf.RngReadAndSkip

      "ResourceScatterMul",
      "ResourceScatterNdAdd",
      "ResourceScatterNdSub",
      "ResourceScatterNdUpdate",
      "ResourceScatterSub",
      "ResourceScatterUpdate",
      "RngReadAndSkip",
      "RngSkip",
      "Roll",
      "ScatterNd",
      "SegmentSumV2",
      "SegmentProdV2",
      "SegmentMinV2",
      "SegmentMaxV2",
      "SelfAdjointEigV2",

  );

  let results = (outs
    TF_FloatTensor:$product
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

def TF_RngReadAndSkipOp : TF_Op<"RngReadAndSkip", []> {
  let summary = "Advance the counter of a counter-based RNG.";

  let description = [{
The state of the RNG after
`rng_read_and_skip(n)` will be the same as that after `uniform([n])`