// CHECK-NEXT: tfrt.return [[out_chain]], [[res]] : !tfrt.chain, !tfrt_fallback.tf_tensor
func.func @gather(%indices: tensor<?xi32>,
             %resource: tensor<*x!tf_type.resource>) -> tensor<*xi32> {
  %0 = "tf.ResourceGather"(%resource, %indices) {batch_dims = 0 : i64, device = "/device:CPU:0", validate_indices = true}: (tensor<*x!tf_type.resource>, tensor<?xi32>) -> (tensor<*xi32>)
  func.return %0 : tensor<*xi32>
}

    %cst = "tf.Const"() <{value = dense<[1]> : tensor<1xi32>}> : () -> tensor<1xi32>
    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "Variable"}> : () -> tensor<!tf_type.resource<tensor<2x3xbf16>>>
    %1 = "tf.ResourceGather"(%0, %cst) <{batch_dims = 0 : i64, validate_indices = true}> : (tensor<!tf_type.resource<tensor<2x3xbf16>>>, tensor<1xi32>) -> tensor<1x3xbf16>
    return %1: tensor<1x3xbf16>
  }

    // CHECK: return [[GATHER]]
    %1 = "tf.ResourceGather"(%resource, %indices) {_xla_outside_compilation = "0"} : (tensor<*x!tf_type.resource<tensor<*xi32>>>, tensor<?xi32>) -> (tensor<*xi32>)
    tf_device.return %1 : tensor<*xi32>
  }) : () -> (tensor<*xi32>)
  func.return %0: tensor<*xi32>
}



// Tests that resource subtype is correctly propagated when decomposing tf.ResourceGather.

  // CHECK: [[var:%.*]] = tf_mlrt.executeop.device([[device]]){{.*}}op: \22ResourceGather\22
  // CHECK: [[rendezvous_key_base:%.*]] = tf_mlrt_tpu.compile_and_execute([[var]])
  %0 = "tf.ResourceGather"(%arg0, %indices) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<4xf32>>>, tensor<i32>) -> tensor<f32>

      ),
      (TF_AssignVariableOp $m_resource, $new_m, (CreateConstBoolAttrFalse)),
      (TF_AssignVariableOp $v_resource, $new_v, (CreateConstBoolAttrFalse))
    ]
  >;

// Pattern to decompose tf.ResourceGather into tf.ReadVariable and tf.GatherV2.
def DecomposeResourceGather : Pat<
  (TF_ResourceGatherOp:$old_result
    $resource, $indices, $batch_dims, $validate_indices),
  (TF_GatherV2Op:$dest

  // 2) Convert composite functions like lstm/rnns, along with proper function
  // inlining & dce.
  // 3) Lower static tensor list pass.

  // This decomposes resource ops like ResourceGather into read-variable op
  // followed by gather. This is used when the saved model import path is used
  // during which resources don't get frozen in the python layer.
  pass_manager->addNestedPass<mlir::func::FuncOp>(

      "ResourceApplyMomentum",
      "ResourceApplyPowerSign",
      "ResourceApplyProximalAdagrad",
      "ResourceApplyProximalGradientDescent",
      "ResourceApplyRMSProp",
      "ResourceGather",
      "ResourceScatterAdd",
      "ResourceScatterDiv",
      "ResourceScatterMax",
      "ResourceScatterMin",
      "ResourceScatterMul",
      "ResourceScatterNdAdd",

    ([CVE-2021-37652](https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2021-37652))
*   Fixes a division by 0 in `ResourceGather`
    ([CVE-2021-37653](https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2021-37653))
*   Fixes a heap OOB and a `CHECK` fail in `ResourceGather`
    ([CVE-2021-37654](https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2021-37654))
*   Fixes a heap OOB in `ResourceScatterUpdate`

  );

  let results = (outs);

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<3>;
}

def TF_ResourceGatherOp : TF_Op<"ResourceGather", []> {
  let summary = [{
Gather slices from the variable pointed to by `resource` according to `indices`.
  }];

  let description = [{