custom_call->setAttr(kShardingAttr, op.get_XlaShardingAttr());
    rewriter.replaceOp(op, custom_call.getResult(0));

    return success();
  }
};

// Converts a TF InplaceUpdate op to DynamicUpdateSlice HLO.
class ConvertInplaceUpdateOp : public OpRewritePattern<TF::InplaceUpdateOp> {
 public:
  using OpRewritePattern::OpRewritePattern;

  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

def TF_InplaceUpdateOp : TF_Op<"InplaceUpdate", [Pure]> {
  let summary = "Updates specified rows 'i' with values 'v'.";

  let description = [{
Computes `x[i, :] = v; return x`.