VpfCdZddlmZddlZddlmZddlmZddlmZddlm Z ddl m Z dd l m Z dd l m Z dd l mZdd lmZdd lmZddlmZddlmZddlmZddlmZ d*dededededededeeeffdZ d*dededededededeeeffdZ dZ!dZ"d Z#d!Z$d"d#d$Z%d%Z&d&Z'ej(d'Z)de)_*e)+eej,e)e)-e!e j.e)ee%d#e j.e)e%d()e&e j/e)<e'e j0e)<dS)+aANN (Approximate Nearest Neighbor) computes top-k with a configurable recall rate. This package only optimizes the TPU backend. For other device types it fallbacks to sort and slice. Usage:: import functools import jax # MIPS := maximal inner product search # Inputs: # qy: f32[qy_size, feature_dim] # db: f32[db_size, feature_dim] # # Returns: # (f32[qy_size, k], i32[qy_size, k]) @functools.partial(jax.jit, static_argnames=["k", "recall_target"]) def mips(qy, db, k=10, recall_target=0.95): dists = jax.lax.dot(qy, db.transpose()) # Computes max_k along the last dimension # returns (f32[qy_size, k], i32[qy_size, k]) return jax.lax.approx_max_k(dists, k=k, recall_target=recall_target) # Multi-core example # Inputs: # qy: f32[num_devices, qy_size, feature_dim] # db: f32[num_devices, per_device_db_size, feature_dim] # db_offset: i32[num_devices] # db_size = num_devices * per_device_db_size # # Returns: # (f32[qy_size, num_devices, k], i32[qy_size, num_devices, k]) @functools.partial( jax.pmap, # static args: db_size, k, recall_target static_broadcasted_argnums=[3, 4, 5], out_axes=(1, 1)) def pmap_mips(qy, db, db_offset, db_size, k, recall_target): dists = jax.lax.dot(qy, db.transpose()) dists, neighbors = jax.lax.approx_max_k( dists, k=k, recall_target=recall_target, reduction_input_size_override=db_size) return (dists, neighbors + db_offset) # i32[qy_size, num_devices, k] pmap_neighbors = pmap_mips(qy, db, db_offset, db_size, 10, 0.95)[1] # i32[qy_size, num_devices * k] neighbors = jax.lax.collapse(pmap_neighbors, start_dimension=1, stop_dimension=3) Todos:: * On host top-k aggregation * Inaccurate but fast differentiation )partialN)ad_util)core)dispatch)dtypes)ad)batching)mlir)xla)lax) xla_client)ir)func)hlo)Arrayffffff?Toperandkreduction_dimension recall_targetreduction_input_size_overrideaggregate_to_topkreturnc Dt||||d||S)aReturns max ``k`` values and their indices of the ``operand`` in an approximate manner. See https://arxiv.org/abs/2206.14286 for the algorithm details. Args: operand : Array to search for max-k. Must be a floating number type. k : Specifies the number of max-k. reduction_dimension : Integer dimension along which to search. Default: -1. recall_target : Recall target for the approximation. reduction_input_size_override : When set to a positive value, it overrides the size determined by ``operand[reduction_dim]`` for evaluating the recall. This option is useful when the given ``operand`` is only a subset of the overall computation in SPMD or distributed pipelines, where the true input size cannot be deferred by the operand shape. aggregate_to_topk : When true, aggregates approximate results to the top-k in sorted order. When false, returns the approximate results unsorted. In this case, the number of the approximate results is implementation defined and is greater or equal to the specified ``k``. Returns: Tuple of two arrays. The arrays are the max ``k`` values and the corresponding indices along the ``reduction_dimension`` of the input ``operand``. The arrays' dimensions are the same as the input ``operand`` except for the ``reduction_dimension``: when ``aggregate_to_topk`` is true, the reduction dimension is ``k``; otherwise, it is greater equals to ``k`` where the size is implementation-defined. We encourage users to wrap ``approx_max_k`` with jit. See the following example for maximal inner production search (MIPS): >>> import functools >>> import jax >>> import numpy as np >>> @functools.partial(jax.jit, static_argnames=["k", "recall_target"]) ... def mips(qy, db, k=10, recall_target=0.95): ... dists = jax.lax.dot(qy, db.transpose()) ... # returns (f32[qy_size, k], i32[qy_size, k]) ... return jax.lax.approx_max_k(dists, k=k, recall_target=recall_target) >>> >>> qy = jax.numpy.array(np.random.rand(50, 64)) >>> db = jax.numpy.array(np.random.rand(1024, 64)) >>> dot_products, neighbors = mips(qy, db, k=10) Trrris_max_krrapprox_top_k_pbindrrrrrrs P/var/www/html/nettyfy-visnx/env/lib/python3.11/site-packages/jax/_src/lax/ann.py approx_max_kr#]s7b     -!$A)  + ++c Dt||||d||S)a Returns min ``k`` values and their indices of the ``operand`` in an approximate manner. See https://arxiv.org/abs/2206.14286 for the algorithm details. Args: operand : Array to search for min-k. Must be a floating number type. k : Specifies the number of min-k. reduction_dimension: Integer dimension along which to search. Default: -1. recall_target: Recall target for the approximation. reduction_input_size_override : When set to a positive value, it overrides the size determined by ``operand[reduction_dim]`` for evaluating the recall. This option is useful when the given operand is only a subset of the overall computation in SPMD or distributed pipelines, where the true input size cannot be deferred by the ``operand`` shape. aggregate_to_topk : When true, aggregates approximate results to the top-k in sorted order. When false, returns the approximate results unsorted. In this case, the number of the approximate results is implementation defined and is greater or equal to the specified ``k``. Returns: Tuple of two arrays. The arrays are the least ``k`` values and the corresponding indices along the ``reduction_dimension`` of the input ``operand``. The arrays' dimensions are the same as the input ``operand`` except for the ``reduction_dimension``: when ``aggregate_to_topk`` is true, the reduction dimension is ``k``; otherwise, it is greater equals to ``k`` where the size is implementation-defined. We encourage users to wrap ``approx_min_k`` with jit. See the following example for nearest neighbor search over the squared l2 distance: >>> import functools >>> import jax >>> import numpy as np >>> @functools.partial(jax.jit, static_argnames=["k", "recall_target"]) ... def l2_ann(qy, db, half_db_norms, k=10, recall_target=0.95): ... dists = half_db_norms - jax.lax.dot(qy, db.transpose()) ... return jax.lax.approx_min_k(dists, k=k, recall_target=recall_target) >>> >>> qy = jax.numpy.array(np.random.rand(50, 64)) >>> db = jax.numpy.array(np.random.rand(1024, 64)) >>> half_db_norm_sq = jax.numpy.linalg.norm(db, axis=1)**2 / 2 >>> dists, neighbors = l2_ann(qy, db, half_db_norm_sq, k=10) In the example above, we compute ``db^2/2 - dot(qy, db^T)`` instead of ``qy^2 - 2 dot(qy, db^T) + db^2`` for performance reason. The former uses less arithmetic and produces the same set of neighbors. Frrr!s r" approx_min_kr&s7j     -!$A)  + ++r$c@|dkrtd|t|jdkr'td|jt |j}|||kr)td|||t j|jtj std||}|r|||<ngtj ||fr;tj|t|||||d||<nt!d|d|d|||j|j ||tjtj fS) Nrzk must be positive, got z5approx_top_k operand must have >= 1 dimension, got {}z;k must be smaller than the size of reduction_dim {}, got {}zoperand must be a floating typezSapprox_top_k with aggregate_to_topk=False not yet implemented when either the `k` (z$) or the reduction dimension size (z) are symbolic)shapedtype weak_type)r(r)) ValueErrorlenr( TypeErrorformatlistr issubdtyper)npfloatingris_constant_shapexcopsApproxTopKReductionOutputSizeNotImplementedErrorupdater*int32) rrrrrrrdimsreduction_input_sizes r"_approx_top_k_abstract_evalr<s!VV 333 4 441 KRR    gm  $ "" ELL $ %q * * + ++  7="+ 6 68 6 7 7712 M !D  3Q788M " D Dc$iiM;L%!'!''(!*D   L L L'; L L L M MM ..t7=#*#4  6 6 ..t28BH+=+=. > > @@r$c tjd|rdnd}tj|dtj|}tj|dtj|}tj|dtjtjtj tj|dtjtjtj |r!tj ||}n tj ||}| |S)Nztop_k_{}_comparatorgtltr)r4 XlaBuilderr.r parameterShape scalar_shaper1r)r9r5GtLtbuild)op_typercp0p1 cmp_results r"_comparator_builderrOsm""8#=44>>@@! }Q280099::" }Q280099::"-1bh++BHRX,>,>??@@@-1bh++BHRX,>,>??@@@ #2r""JJ2r""J   r$c`tj|r tj n tj|S)N)r))r1arrayinf)rJrs r"_get_init_val_literalrSs& X126''26 A A AAr$ctjg|}tjgtjd}||||g}tjgtjdg}tj||}tj|jj j 5tj d |rdnd||}dddn #1swxYwY|jj||} tj| 5| j\} } } } t$j|rdnd} t%j| | | }t%j|gdddn #1swxYwY|S) N r@ztop_k_{}_{}_comparatorr>r?GTLT)comparison_direction)rRankedTensorTypeget IntegerType get_signless FunctionTypeInsertionPointat_block_beginmodule_contextmodulebodyrFuncOpr. symbol_tableinsertadd_entry_block argumentsrComparisonDirectionAttrcomparereturn_)ctxrJrscalarindexir_types result_typescomparator_type comparator entry_blockrLrM_ directionrNs r"_comparator_builder_mlirru s)  " "2w / /&  ! !"bn&A&A"&E&E F F%feU +(%))"bn.I.I!.L.LMMN,O'',??/ ''(:(A(FGG ''(BdGLLJ!((444**,,+ %%(LBAq+//0JdKKIR)DDDJK   s%0.D**D.1D.z)_approx_top_k_lowering..$s/ C CZ4+ , , C C C C C Cr$rz"operand must be an array, but was ) dimensionr) reduction_dimrrr is_fallbackc3HK|]}tj|jVdSry)rr3r()r|aval_outs r"r~z)_approx_top_k_lowering..As/NNH  / /NNNNNNr$chg|].}tjtj|j/Sr)r shape_tensoreval_dynamic_shaper()r|rrks r" z*_approx_top_k_lowering..DsC'''  $1#x~FFGG'''r$top_k ApproxTopKc6g|]}tj|Srr aval_to_ir_typer|avals r"rz*_approx_top_k_lowering..L#KKKTd*400KKKr$)rooperandscalled_computationsbackend_config result_shapeszstablehlo.dynamic_approx_top_kc6g|]}tj|Srrrs r"rz*_approx_top_k_lowering..Urr$)!avals_inallr(r,r+r dtype_to_ir_typer)rF32TyperZruiotarr{r1r9rconstantDenseElementsAttrrS ir_constantreshapei64_attr FloatAttrBoolAttr avals_outis_constant_dim custom_callnamevalueeval_dynamic_shape_as_valsresults)rkrrrrrrrrvop_shapeop_dimsrJ recall_typerqrinit_arginit_val_arrayinit_valrroutk_values` r"_approx_top_k_loweringrs  C Ccl C C C C CCCC \!_ "(]]a D(DD E EE '  !#,q/"7 8 8'   +1g,,)<<'Wh??* 3(a)>II0 2 2 2$\".2228B<<@@ A A((a)>II.  n44R88 9 9(m$788g'++KGG'*../@AA m122 44. C$(G$4$8$8$B$BN=!NN NNNNN'MM'''' '''M !%"mA..N7  KKS]KKK484'_23%#  % % %CC.sQD99HG  (KKS]KKK48W='_23%#  % % %C r$c  t|dksJt|dksJ|\}|\ fdt|jD} | |}t||||||| ffS)Nr@cg|]}|u| Srr)r|d batch_axiss r"rz,_approx_top_k_batch_rule..ds# C C C1q /B/BQ/B/B/Br$r)r,rangendimrr ) batch_operands batch_axesrrrrrrrdim_maprs @r"_approx_top_k_batch_ruler]s ^   ! ! ! ! ZA     ('+* C C C Cgl++ C C C' 34     -!$A)  + +.8,D EEr$c|\}|\} |rt|||||\} nt|||||\} t| tjur tj| } nojt} dkr| z fdt| Dtfdt| D} | | } | f| tjffS)NrcFg|]}tjj|Sr)r broadcasted_iotar))r|iarg_out arg_shapes r"rz%_approx_top_k_jvp..s7   >?W]Iq99   r$c38K|]}|krn|VdSryr)r|rriotasrs r"r~z$_approx_top_k_jvp..sOPP>?1+++qPPPPPPr$) r#r&typerZero from_valuer(r,rtuple)primalstangentsrrrrrrrtangentval_out tangent_outrankidxrrrs ` @@@r"_approx_top_k_jvprxsq(' (' 7#GQ0C$1$A$577GWW $GQ0C$1$A$577GW ']]gl"",))'22KK I y>>DQT!     CH;;   E PPPPPPCH;;PPP P PC#,K 7 k7<+B+B7+K+KL LLr$ approx_top_ktpu)platform)rrrT)1__doc__ functoolsrnumpyr1jax._srcrrrrjax._src.interpretersrr r r jax._src.laxr jax._src.libr r4jax._src.lib.mlirrjax._src.lib.mlir.dialectsrrjax._src.typingrintfloatboolrr#r&r<rOrSrurrr Primitivermultiple_resultsdef_implapply_primitivedef_abstract_evalregister_loweringprimitive_batchersprimitive_jvpsrr$r"rs[77r$$$$$$******&&&&&&%%%%%%)))))) ++++++******!!!!!! -/(,68+/ 8+8+%8+8+&)8+!&8+14 8+ %) 8+ 5:%,4G 8+8+8+8+z-/(,68+/ <+<+%<+<+&)<+!&<+14 <+ %) <+ 5:%,4G <+<+<+<+~@@@B   BBB2?D<<<<<|EEE6MMM> //"& 8.IIJJJ !<===~g4tDDDFFF~'=!&((((.FN+$5.!!!r$