}
  }
  return resource_vec;
}

LogicalResult _TfrtGetResourceOp::verify() {
  _TfrtGetResourceOp get_resource_op = *this;
  // The sizes of indices, shared_name and container must be equal.
  int32_t indices_size =
      get_resource_op->getAttrOfType<mlir::ArrayAttr>("indices").size();
  int32_t shared_name_size =
      get_resource_op->getAttrOfType<mlir::ArrayAttr>("shared_name").size();
  int32_t container_size =

// Gathers elements in buffer with the indices.
Value GatherElements(Value indices, Value buffer, OpBuilder builder,
                     Location loc);

// Scatters elements into buffer, where each scattered element is accumulated
// with the old value in buffer.
Value ScatterAccumulateElements(Value indices, Value updates, Value buffer,

      /*methodBody=*/[{
        return $_op.getDataFormat();
    }]>,
    InterfaceMethod<
      [{Returns indices of layout dependent arguments.}],
      "SmallVector<unsigned, 4>", "GetLayoutDependentArgs", (ins)
    >,
    InterfaceMethod<
      [{Returns indices of layout dependent results.}],
      "SmallVector<unsigned, 4>", "GetLayoutDependentResults", (ins)
    >,
    InterfaceMethod<

                      ConversionPatternRewriter& rewriter);

// If index_vector_dim == indices.rank() then insert the implicit extra
// dimension into indices to normalize everything to index_vector_dim ==
// indices.rank() - 1.
LogicalResult NormalizeIndexVector(Operation* parent_op, Value& indices,
                                   ShapedType& indices_type,
                                   int64_t index_vector_dim,

#endif

typedef struct TF_VariableInputLockHolder TF_VariableInputLockHolder;

// Expose higher level Assignment operation for Pluggable vendors to implement
// in the plugin for Training. The API takes in the context with indices for
// the input and value tensors. It also accepts the copy callback provided by
// pluggable vendor to do the copying of the tensors. The caller takes ownership

  require(name.isNotEmpty()) { "name is empty" }
  for (i in name.indices) {
    val c = name[i]
    require(c in '\u0021'..'\u007e') {
      "Unexpected char 0x${c.charCode()} at $i in header name: $name"
    }
  }
}

internal fun headersCheckValue(
  value: String,
  name: String,
) {
  for (i in value.indices) {
    val c = value[i]
    require(c == '\t' || c in '\u0020'..'\u007e') {

// initialized, the value will not be present.
using ResourceVarsSnapshot = absl::flat_hash_map<int, std::optional<Tensor>>;

// Takes a snapshot of the values of resource variable arguments, whose indices
// are specified in `variable_indices` argument. We snapshot tensors that back
// resource variables since concurrent updates may modify the shape, and it is

  if (!ranked_ty) return UnrankedTensorType::get(element_ty);
  int64_t rank = ranked_ty.getRank();

  DenseIntElementsAttr indices;
  if (!matchPattern(reduction_indices, m_Constant(&indices))) {
    // Output type is unranked if reduction indices are not constant and reduced
    // dimensions are not kept.
    if (!keep_dims.getValue()) return UnrankedTensorType::get(element_ty);

  [(IsNotInLiftedFunc $res), (IsConstTensor $b)], [], (addBenefit 1)>;

def LiftGather : Pat<
  (TF_GatherV2Op:$res $params, $indices, $axis, $batch_dims),
  (LiftAsTFPartitionedCall<"composite_gather_fn">
    (ArgumentList $params, $indices, $axis),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"batch_dims"> $batch_dims))),

    $constant_operand_indices are the indices of the inputs that are constant to the TPU program (e.g. weights in inference), the rest of the inputs are input tensors.
    constant_operand_indices is sorted in ascending order.
    $operands_with_static_shape are indices of operands that are tagged with a maximum static shape.