ukiP  ddlmZddlZddlmZmZddlZddlZddlm Z m Z m Z m Z m Z ddlZddlmZddlmZddlmZddlmZdd lmZmZmZmZmZdd lmZdd lmZdd lmZdd lm Z!ddlm"Z#ddl$m%Z%m&Z&m'Z'm(Z(m)Z)ddl*m+Z+m,Z,m-Z-m.Z.e.dZ/ej`Z1e e2ee2dfZ3dwdZ4dxdZ5dydZ6dzdZ7d{dZ8ee e ge fZ9ddddddddddddd d|dZ:dZ;d}dZddZ?dd Z@dd!ZAejd"d# d dd$ZCe/ d dd%ZDejd"d# d dd&ZEe/ d dd'ZFejd(d# d dd)ZGe/ d dd*ZHejd+d# d dd,ZIe/ d dd-ZJejd(d# ddd. dd/ZKe/ ddd. dd0Z0ejd(d# ddd. dd1ZLe/ ddd. dd2ZMejd+d# d dd3ZNejd+d# d dd4ZOejd+d# d dd5ZPejd+d# d dd6ZQejd+d# d dd7ZRejd+d# d dd8ZS d dd9ZT d dd:ZUe/ d dd;ZVe/ d dd<ZWdd=ZXe/ ddd. dd>ZY dd?ZZejd@d# ddddA ddBZ[e d ddCZ\e dddD ddEZ\e d ddFZ\e/ d ddGZ\ejdHd# d ddIZ]e/ dddddJ ddKZ^ejdLMddddddddN ddOZ_ddPZ`e/ dddddJ ddQZaejdLMddddddddR ddSZbe/ d ddTZcejd(M d ddUZde/ejd(M d ddVZe d ddWZfe/ejd(M d ddXZge/ejd(M d ddYZhe/ejdLM d ddZZie/ejdLM d dd[Zje/ejdLM d dd\Zke/ejdLM d dd]Zldddddddd^ dd_Zme/ejdLM d dd`ZnGdadbe Zo d ddcZpe/ejddM d ddeZqe/ejddM d ddfZre/ejddM d ddgZse/ejddM d ddhZtejddM d ddiZue/ddddj ddkZve/ddddj ddlZwe/ejdmM de)dn ddoZxe/ejdmM de)dn ddpZy ddqZze/ejdmM de)dn ddrZ{e/ejdmM de)dn ddsZ|e/ejdtM d dduZ}e/ejdtM d ddvZ~y)) annotationsN)CallableSequence)overloadAnyLiteralProtocolUnion)api)config)core)dtypes) _broadcast_toensure_arraylikepromote_dtypes_inexactpromote_dtypes_numeric_where) control_flow)lax)other)parallel)slicing)Array ArrayLikeDType DTypeLike DeprecatedArg)canonicalize_axiscanonicalize_axis_tuplemaybe_named_axis set_modulez jax.numpyct|dr|j}tj|xst j |S)N __jax_array__)hasattrr#ris_python_scalarnpisscalar)elements T/home/cdr/jupyterlab/.venv/lib/python3.12/site-packages/jax/_src/numpy/reductions.py _isscalarr*0s9 Wo&##%G   ) AR[[-AAcRtd|}t|tj|}t|tj|}t tj|Dcgc] }||k7s | }}|j ||t j||Scc}w)Nmoveaxis)rrr&ndimrangeinsertr transpose)asource destinationiperms r) _moveaxisr75sz1%! VRWWQZ 0&!+rwwqz:+2771:& 6!v+! 6$ 6++k6" q$  7s + B$6B$ctj|tjtjfvrtjdStj|S)Nfloat32)r&dtypefloat16rbfloat16r:s r) _upcast_f16r>>s<XXe_V__55 88I  %r+c^tj|tjrtjStj|tj rWtj }tj|jtj|jkr|S|Stj|tjrUtj}tj|jtj|jkr|S|SN) r issubdtyper&bool_default_int_dtypeunsignedintegerdefault_uint_dtypeiinfobitsinteger)r:rErCs r)_promote_integer_dtyperICs ubhh'  # # %%  2 23224 xxbhh'9:???  ,  +002 xxbhh'89>>>  ,r+c||St||}|jtk7rtd|d|jd|S)Nzjnp.z9: where must be None or a boolean array; got where.dtype=.)rr:bool ValueError)namewheres r) check_whererPRsO ] L 4 '% [[D  TFL ~QO  ,r+TF) has_identitypreprocbool_opupcast_f16_for_computationaxisr:outkeepdimsinitialwhere_parallel_reducepromote_integersc h|xs|}| td|dt||}t|| } tjd|d|d}| |s| t d|d|r||n|}t ||\}}| :|s8tj|tfd|Dst d|d | |j}|r"t|}ntj| |}|r0tj|tjr t!|}n|}t#j$||}|tj&k7r|n|}t)||}| t+| ||}||ur|td |d |||}nt#j,||||}| ct#j$| t#j.|j}|jd k7rt d |j|||}| rt#j0||}t#j$|| xs|S)NThe 'out' argument to jnp.z is not supported.zaxis argument to jnp.z().zreduction operation zO does not have an identity, so to use a where mask one has to specify 'initial'c3.K|] }|dk\ywN).0dshapes r) z_reduction..{s0!aA 0sz'zero-size array to reduction operation z which has no identityz)Named reductions not implemented for jnp.z()raz3initial value must be a scalar. Got array of shape )NotImplementedErrorrrPr concrete_or_errorrM_reduction_dimsr&rd_allr:rIr!check_and_canonicalize_user_dtyperAinexactr>rconvert_element_typerB_reduction_init_valrreduceasarray expand_dims)r2rNopinit_valrQrRrSrTrUr:rVrWrXrYrZr[pos_dimsdims result_dtypecomputation_dtyperesult initial_arrrds @r) _reductionry^sW Mr' _  :4&@RS TTtQ! tV $&  d.CD6,M N$ _\f.@ +D62?@ AAgaj!"1d+.(D _\ HHQKE 0x0 0 @F\] ^^ ]77L+L9l;;E4HLF$5$5lBJJ$O#L1$ q"34! "(( *r"!H -( vq(#A T "KD6QS T UU Q %F ZZ8R .F **7CKKN4H4HIKB --8->->,?A BB  V $F __VX .F ! !&%*?< @@r+c&t|fddS)Nct|Sr@r)r5ranks r)z0_canonicalize_axis_allow_named..s'8D'Ar+c|Sr@ra)rNs r)r~z0_canonicalize_axis_allow_named..sPTr+)r )xr}s `r)_canonicalize_axis_allow_namedrs !ACT UUr+c|+tttjfdzSt |tj tt fs|f}tfd|D}t|tt|k7rtd|td|D}t|t|k7r||fS||fS)Nc3\K|]#}t|tj%ywr@)rr&r.)rbrr2s r)rez"_reduction_dims..s($4ArwwqzB$s),zduplicate value in 'axis': c3BK|]}t|ts|ywr@) isinstanceint)rbrs r)rez"_reduction_dims..sEq*Q2DEs) tupler/r&r.rndarraylistlensetrM)r2rU canon_axiscanon_pos_axiss` r)rhrhs \ % # $ & ** dRZZ5 6 7D$"$$*_C O,, 24&9 ::EJEE.C O+ : %% z !!r+c|j}|dk(rtj|dkD|Stj|rt j |tj r}t j|shtjtj|rt j|jnt j|j|} tj||S#t$rt j |tjsJtj|t j|}}tj|dkr |jn |j|cYSwxYw)NrLrr=)r:r&arrayisinfrrAfloating supports_infisneginffinfominmax OverflowErrorrHsignrF)r2rra_dtyperinfos r)rmrms GG'  88HqL 00hhxV..w D!!'*xxR[[5J W-11#\\'266gGHG 88HG ,, G   Wbjj 11 1"FLL$9$D 88qDHHdhhg FFGs#C::B FFctj|jtjr |j }t j|tjSr@) rrAr:r&complexfloatingrealrrlrBoperands r) _cast_to_boolrs< w}}b&8&89llG ! !'288 44r+ct|dS)Nr)rrs r)_cast_to_numericrs  ( ++r+cvtj|jdstd|j|S)N)rLintegralz%integer argument required; got dtype=)risdtyper:rM)arrs r)_require_integerrs1  #7 8 .force..s'Ma 1c*1q0AAMs13)rr TypeErrorr)rs r)forcez$_ensure_optional_axes..forcesAy N ^^A  N M1M MMNs 88z+The axis argument must be known statically.)r rg)rrs r)_ensure_optional_axesrs%N    1; ==r+)rUr:rWr[)static_argnamesinlinect|dtjdttj|du||||||t j |S)Nsumr) rRrSrTrUr:rVrWrXrYrZr[)ryraddr bitwise_or lax_parallelpsumr2rUr:rVrWrXrOr[s r) _reduce_sumrsG Aucggq2BNNQU Uh#E>> x = jnp.array([[1, 3, 4, 2], ... [5, 2, 6, 3], ... [8, 1, 3, 9]]) >>> jnp.sum(x) Array(47, dtype=int32) If ``axis=1``, the sum is computed along axis 1. >>> jnp.sum(x, axis=1) Array([10, 16, 21], dtype=int32) If ``keepdims=True``, ``ndim`` of the output is equal to that of the input. >>> jnp.sum(x, axis=1, keepdims=True) Array([[10], [16], [21]], dtype=int32) To include only specific elements in the sum, you can use ``where``. >>> where=jnp.array([[0, 0, 1, 0], ... [0, 0, 1, 1], ... [1, 1, 1, 0]], dtype=bool) >>> jnp.sum(x, axis=1, keepdims=True, where=where) Array([[ 4], [ 9], [12]], dtype=int32) >>> where=jnp.array([[False], ... [False], ... [False]]) >>> jnp.sum(x, axis=0, keepdims=True, where=where) Array([[0, 0, 0, 0]], dtype=int32) rUr:rVrWrXrOr[)rrrs r)rrs*L Q2483&u&6 88r+cxt|dtjdttj|du|||||||S)Nprodr`) rRrSrTrUr:rVrWrXrYr[)ryrmulr bitwise_andrs r) _reduce_prodr5sA Avsww3COORVUh#EDT VVr+c :t|t|||||||S)aD Return product of the array elements over a given axis. JAX implementation of :func:`numpy.prod`. Args: a: Input array. axis: int or array, default=None. Axis along which the product to be computed. If None, the product is computed along all the axes. dtype: The type of the output array. Default=None. keepdims: bool, default=False. If true, reduced axes are left in the result with size 1. initial: int or array, Default=None. Initial value for the product. where: int or array, default=None. The elements to be used in the product. Array should be broadcast compatible to the input. promote_integers : bool, default=True. If True, then integer inputs will be promoted to the widest available integer dtype, following numpy's behavior. If False, the result will have the same dtype as the input. ``promote_integers`` is ignored if ``dtype`` is specified. out: Unused by JAX. Returns: An array of the product along the given axis. See also: - :func:`jax.numpy.sum`: Compute the sum of array elements over a given axis. - :func:`jax.numpy.max`: Compute the maximum of array elements over given axis. - :func:`jax.numpy.min`: Compute the minimum of array elements over given axis. Examples: By default, ``jnp.prod`` computes along all the axes. >>> x = jnp.array([[1, 3, 4, 2], ... [5, 2, 1, 3], ... [2, 1, 3, 1]]) >>> jnp.prod(x) Array(4320, dtype=int32) If ``axis=1``, product is computed along axis 1. >>> jnp.prod(x, axis=1) Array([24, 30, 6], dtype=int32) If ``keepdims=True``, ``ndim`` of the output is equal to that of the input. >>> jnp.prod(x, axis=1, keepdims=True) Array([[24], [30], [ 6]], dtype=int32) To include only specific elements in the sum, you can use a``where``. >>> where=jnp.array([[1, 0, 1, 0], ... [0, 0, 1, 1], ... [1, 1, 1, 0]], dtype=bool) >>> jnp.prod(x, axis=1, keepdims=True, where=where) Array([[4], [3], [6]], dtype=int32) >>> where = jnp.array([[False], ... [False], ... [False]]) >>> jnp.prod(x, axis=1, keepdims=True, where=where) Array([[1], [1], [1]], dtype=int32) r)rrrs r)rr@s*N a3D9''7 99r+rUrWct|dtjtj d||||||t j  S)NrFrQrUr:rVrWrXrYrZ)ryrrr&infrpmaxr2rUr:rVrWrXrOs r) _reduce_maxrs= AucggwUUh#E>> x = jnp.array([[9, 3, 4, 5], ... [5, 2, 7, 4], ... [8, 1, 3, 6]]) >>> jnp.max(x) Array(9, dtype=int32) If ``axis=1``, the maximum will be computed along axis 1. >>> jnp.max(x, axis=1) Array([9, 7, 8], dtype=int32) If ``keepdims=True``, ``ndim`` of the output will be same of that of the input. >>> jnp.max(x, axis=1, keepdims=True) Array([[9], [7], [8]], dtype=int32) To include only specific elements in computing the maximum, you can use ``where``. It can either have same dimension as input >>> where=jnp.array([[0, 0, 1, 0], ... [0, 0, 1, 1], ... [1, 1, 1, 0]], dtype=bool) >>> jnp.max(x, axis=1, keepdims=True, initial=0, where=where) Array([[4], [7], [8]], dtype=int32) or must be broadcast compatible with input. >>> where = jnp.array([[False], ... [False], ... [False]]) >>> jnp.max(x, axis=0, keepdims=True, initial=0, where=where) Array([[0, 0, 0, 0]], dtype=int32) rUrVrWrXrO)rrr2rUrVrWrXrOs r)rrs'N Q248c&u FFr+)rUrWr:ct|dtjtjd||||||t j  S)NrFr)ryrrr&rrpminrs r) _reduce_minrs; AucggrvvEUh#E>> x = jnp.array([[2, 5, 1, 6], ... [3, -7, -2, 4], ... [8, -4, 1, -3]]) >>> jnp.min(x) Array(-7, dtype=int32) If ``axis=1``, the minimum is computed along axis 1. >>> jnp.min(x, axis=1) Array([ 1, -7, -4], dtype=int32) If ``keepdims=True``, ``ndim`` of the output will be same of that of the input. >>> jnp.min(x, axis=1, keepdims=True) Array([[ 1], [-7], [-4]], dtype=int32) To include only specific elements in computing the minimum, you can use ``where``. ``where`` can either have same dimension as input. >>> where=jnp.array([[1, 0, 1, 0], ... [0, 0, 1, 1], ... [1, 1, 1, 0]], dtype=bool) >>> jnp.min(x, axis=1, keepdims=True, initial=0, where=where) Array([[ 0], [-2], [-4]], dtype=int32) or must be broadcast compatible with input. >>> where = jnp.array([[False], ... [False], ... [False]]) >>> jnp.min(x, axis=0, keepdims=True, initial=0, where=where) Array([[0, 0, 0, 0]], dtype=int32) r)rrrs r)rrs'L Q248c&u FFr+rOc Nt|dtjdt|||| S)NallTrRrUrVrWrY)ryrrrr2rUrVrWrOs r) _reduce_allr1s* Aucoot]3% IIr+c4t|t||||S)aTest whether all array elements along a given axis evaluate to True. JAX implementation of :func:`numpy.all`. Args: a: Input array. axis: int or array, default=None. Axis along which to be tested. If None, tests along all the axes. keepdims: bool, default=False. If true, reduced axes are left in the result with size 1. where: int or array of boolean dtype, default=None. The elements to be used in the test. Array should be broadcast compatible to the input. out: Unused by JAX. Returns: An array of boolean values. Examples: By default, ``jnp.all`` tests for True values along all the axes. >>> x = jnp.array([[True, True, True, False], ... [True, False, True, False], ... [True, True, False, False]]) >>> jnp.all(x) Array(False, dtype=bool) If ``axis=0``, tests for True values along axis 0. >>> jnp.all(x, axis=0) Array([ True, False, False, False], dtype=bool) If ``keepdims=True``, ``ndim`` of the output will be same of that of the input. >>> jnp.all(x, axis=0, keepdims=True) Array([[ True, False, False, False]], dtype=bool) To include specific elements in testing for True values, you can use a``where``. >>> where=jnp.array([[1, 0, 1, 0], ... [0, 0, 1, 1], ... [1, 1, 1, 0]], dtype=bool) >>> jnp.all(x, axis=0, keepdims=True, where=where) Array([[ True, True, False, False]], dtype=bool) rUrVrWrO)rrrs r)rr8#^ Q248c&e 55r+c Nt|dtjdt|||| S)NanyFr)ryrrrrs r) _reduce_anyrjs* Aucnne]3% IIr+c4t|t||||S)aTest whether any of the array elements along a given axis evaluate to True. JAX implementation of :func:`numpy.any`. Args: a: Input array. axis: int or array, default=None. Axis along which to be tested. If None, tests along all the axes. keepdims: bool, default=False. If true, reduced axes are left in the result with size 1. where: int or array of boolean dtype, default=None. The elements to be used in the test. Array should be broadcast compatible to the input. out: Unused by JAX. Returns: An array of boolean values. Examples: By default, ``jnp.any`` tests along all the axes. >>> x = jnp.array([[True, True, True, False], ... [True, False, True, False], ... [True, True, False, False]]) >>> jnp.any(x) Array(True, dtype=bool) If ``axis=0``, tests along axis 0. >>> jnp.any(x, axis=0) Array([ True, True, True, False], dtype=bool) If ``keepdims=True``, ``ndim`` of the output will be same of that of the input. >>> jnp.any(x, axis=0, keepdims=True) Array([[ True, True, True, False]], dtype=bool) To include specific elements in testing for True values, you can use a``where``. >>> where=jnp.array([[1, 0, 1, 0], ... [0, 1, 0, 1], ... [1, 0, 1, 0]], dtype=bool) >>> jnp.any(x, axis=0, keepdims=True, where=where) Array([[ True, False, True, False]], dtype=bool) r)rrrs r)rrqrr+c tj|}tjdj |xs |j }t |dtj|tt|||||| S)Nreduce_bitwise_and rNrqrrrRrUr:rVrWrXrY) rror&rastyper:ryrrr) r2rUr:rVrWrXrOrrrs r)_reduce_bitwise_andrsd A# XXb\ !3#)) 4( C2sQYcs.t4EsU]#E 33r+c dt|dtjdtt |||||| S)Nreduce_bitwise_orrr)ryrrrrrs r)_reduce_bitwise_orrs5 A/CNNQXh.t4EsU]#E 33r+c dt|dtjdtt |||||| S)Nreduce_bitwise_xorrr)ryr bitwise_xorrrrs r)_reduce_bitwise_xorrs5 A0S__qZj.t4EsU]#E 33r+c dt|dtjdtt |||||| S)Nreduce_logical_andTr)ryrrrrrs r)_reduce_logical_andrs5 A0S__t]j.t4EsU]#E 33r+c dt|dtjdtt |||||| S)Nreduce_logical_orFr)ryrrrrrs r)_reduce_logical_orrs5 A/CNNU\i.t4EsU]#E 33r+c dt|dtjdtt |||||| S)Nreduce_logical_xorFr)ryrrrrrs r)_reduce_logical_xorrs5 A0S__u^k.t4EsU]#E 33r+c B| td|tj|d}~td|}t d|}t |\}t ||\}} t|j| ||tj } tjtjtj| | tj| d} |r| ntj | |} tj"tj$|| j'|j(} | j+| ||} tj,tj.| | j'| j(}||St1j2||S)z7Compute log(sum(exp(a))) while avoiding precision loss.z&?L (;(;GV(LLr+c| td|tj|d}td|}t d|}t t jd}|||z}t||z||||||z S)z(Compute log2(sum(2 ** a)) via logsumexp.z=The 'out' argument to jnp.logaddexp2.reduce is not supported.zjnp.logaddexp2.reduce logsumexp2r)rUr:rWrOrX) rfrrjrrPfloatr&rr)r2rUr:rVrWrXrOln2s r) _logsumexp2r s _ ] ^^   4 4 & (E|Q'! lE *% bffQi#  sNG AG$eh 247 88r+c$t||||||S)zAlias of :func:`jax.numpy.min`.r)rrs r)aminr  QTsXE ++r+c$t||||||S)zAlias of :func:`jax.numpy.max`.r)rrs r)rr rr+ct|ttfs|f}n|}d}tj||D]#}|t |fdt jz}%|S)Nr`c|Sr@ra)r5a_shapes r)r~z_axis_size..s '!*r+)rrrr&rdr r axis_size)r2rUaxis_seqsizers @r) _axis_sizers` D5$- (#gHH $ HHQK' Na Q 4l6L6L MMDN +r+c <t|t||||||duS)aReturn the mean of array elements along a given axis. JAX implementation of :func:`numpy.mean`. Args: a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the mean to be computed. If None, mean is computed along all the axes. dtype: The type of the output array. If None (default) then the output dtype will be match the input dtype for floating point inputs, or be set to float32 or float64 for non-floating-point inputs. keepdims: bool, default=False. If true, reduced axes are left in the result with size 1. where: optional, boolean array, default=None. The elements to be used in the mean. Array should be broadcast compatible to the input. out: Unused by JAX. Returns: An array of the mean along the given axis. Notes: For inputs of type `float16` or `bfloat16`, the reductions will be performed at float32 precision. See also: - :func:`jax.numpy.average`: Compute the weighted average of array elements - :func:`jax.numpy.sum`: Compute the sum of array elements. Examples: By default, the mean is computed along all the axes. >>> x = jnp.array([[1, 3, 4, 2], ... [5, 2, 6, 3], ... [8, 1, 2, 9]]) >>> jnp.mean(x) Array(3.8333335, dtype=float32) If ``axis=1``, the mean is computed along axis 1. >>> jnp.mean(x, axis=1) Array([2.5, 4. , 5. ], dtype=float32) If ``keepdims=True``, ``ndim`` of the output is equal to that of the input. >>> jnp.mean(x, axis=1, keepdims=True) Array([[2.5], [4. ], [5. ]], dtype=float32) To use only specific elements of ``x`` to compute the mean, you can use ``where``. >>> where = jnp.array([[1, 0, 1, 0], ... [0, 1, 0, 1], ... [1, 1, 0, 1]], dtype=bool) >>> jnp.mean(x, axis=1, keepdims=True, where=where) Array([[2.5], [2.5], [6. ]], dtype=float32) N)rOrT)_meanr)r2rUr:rVrWrOs r)meanr s-@ q'-uc8  HHr+c&|b|)tjtj|}ntjt ||}t j ||}|Stt|tj||||}|S)N)r:rW) r dimension_as_valuer&rrrrlrrrd)r2rUrWrOr:counts r)_countrcs| ] |%%bggaj1e%%jD&9:e  $ $UE 2E ,  eRXXa[14ux XE ,r+)rUr:rWrT)rTrOc td|}td|}| td| tj|j }ntj |d}|r0tj|tjr t|}n|}t|||||} tjt|||||| j|S)Nrz0The 'out' argument to jnp.mean is not supported.rUrWrOr:r:rWrO)rrPrfrto_inexact_dtyper:rjrAr&rkr>rrdivrr) r2rUr:rVrWrTrOrurv normalizers r)rrts vq!! fe $%_ P QQ ]**1773L;;E6JLF$5$5lBJJ$O#L1$    *  !T*XUK  F<r+cyr@rar2rUweightsreturnedrWs r)averager(sPSr+)rWcyr@rar%s r)r(r(sGJr+cyr@rar%s r)r(r(s\_r+c2t|t||||S)a5Compute the weighed average. JAX Implementation of :func:`numpy.average`. Args: a: array to be averaged axis: an optional integer or sequence of integers specifying the axis along which the mean to be computed. If not specified, mean is computed along all the axes. weights: an optional array of weights for a weighted average. This must either exactly match the shape of `a`, or if `axis` is specified, it must have shape ``a.shape[axis]`` for a single axis, or shape ``tuple(a.shape[ax] for ax in axis)`` for multiple axes. returned: If False (default) then return only the average. If True then return both the average and the normalization factor (i.e. the sum of weights). keepdims: If True, reduced axes are left in the result with size 1. If False (default) then reduced axes are squeezed out. Returns: An array ``average`` or tuple of arrays ``(average, normalization)`` if ``returned`` is True. See also: - :func:`jax.numpy.mean`: unweighted mean. Examples: Simple average: >>> x = jnp.array([1, 2, 3, 2, 4]) >>> jnp.average(x) Array(2.4, dtype=float32) Weighted average: >>> weights = jnp.array([2, 1, 3, 2, 2]) >>> jnp.average(x, weights=weights) Array(2.5, dtype=float32) Use ``returned=True`` to optionally return the normalization, i.e. the sum of weights: >>> jnp.average(x, returned=True) (Array(2.4, dtype=float32), Array(5., dtype=float32)) >>> jnp.average(x, weights=weights, returned=True) (Array(2.5, dtype=float32), Array(10., dtype=float32)) Weighted average along a specified axis: >>> x = jnp.array([[8, 2, 7], ... [3, 6, 4]]) >>> weights = jnp.array([1, 2, 3]) >>> jnp.average(x, weights=weights, axis=1) Array([5.5, 4.5], dtype=float32) )_averagerr%s r)r(r(sn !*40'8X NNr+)rUr'rWc dnttj|tdt \t |}At jdtjj|j}n>ttr-t jdtjfdD|j}ntd|\}t |\}j |j k7r t#d|j tfdDk7r t#dtfd t%j D}t j&||ttj( }t+||}t+|z||z }|r3|j j k7rt-||j }||fS|S) Nr(rrar=c3bK|]&}tjj|(ywr@)r rrd)rbrcr2s r)rez_average..s%*]ST4+B+B1771:+N*]s,/z;Axis must be specified when shapes of a and weights differ.c3<K|]}j|ywr@)rd)rbaxr2s r)rez_average..s; ;szIShape of weights must be consistent with shape of a along specified axis.c34K|]\}}|vr|ndywr_ra)rbr5dimrUs r)rez_average..s QFAsqDya/Qs) dimensions)rr&r.rrrrfullr rrr:rrmathrrdrM enumeratereshapeargsortrr)r2rUr&r'rWavg weights_sum new_shapes`` r)r,r,s#:4#L$ _A&A  "BA qth /C |HHR!8!8!@ Rk D% HHR*]X\*]!]ehenenok!)Q8JAw'73JAwww'-- +, , %;d;; ;67 7Qi>PQQi GY5DAQ;RSggD8>> x = jnp.array([[1, 3, 4, 2], ... [5, 2, 6, 3], ... [8, 4, 2, 9]]) >>> with jnp.printoptions(precision=2, suppress=True): ... jnp.var(x) Array(5.74, dtype=float32) If ``axis=1``, variance is computed along axis 1. >>> jnp.var(x, axis=1) Array([1.25 , 2.5 , 8.1875], dtype=float32) To preserve the dimensions of input, you can set ``keepdims=True``. >>> jnp.var(x, axis=1, keepdims=True) Array([[1.25 ], [2.5 ], [8.1875]], dtype=float32) If ``ddof=1``: >>> with jnp.printoptions(precision=2, suppress=True): ... print(jnp.var(x, axis=1, keepdims=True, ddof=1)) [[ 1.67] [ 3.33] [10.92]] To include specific elements of the array to compute variance, you can use ``where``. >>> where = jnp.array([[1, 0, 1, 0], ... [0, 1, 1, 0], ... [1, 1, 1, 0]], dtype=bool) >>> with jnp.printoptions(precision=2, suppress=True): ... print(jnp.var(x, axis=1, keepdims=True, where=where)) [[2.25] [4. ] [6.22]] r5ddof and correction can't be provided simultaneously.varrUr:rVr<rWrOa_mean)rrrMr_varr r2rUr:rVddofrWrOrr<s r)r?r?s^fJ dC DAI L MMua ! a+D1CT^iq$ ((r+rUr:rW)rr@ctd|}|tj|d}| tdt |j |\}}t j|j|}|t|||d|}ntd|j|}t j||} tj|tjrIt jt j | t j"| } | j }nt j$| } t'|||||} t j| t j(||} t+| ||||} t j,| | j|} t/j0d5t3| dkD| tj4} ddd| S#1swY| SxYw)Nr?z/The 'out' argument to jnp.var is not supported.Tr rFr)rPrrjrf_var_promote_typesr:rrorrrrrAr&rrrconjsquarerrlrr"r debug_nansrnan) r2rUr:rVr<rWrOrArvcenteredr#rws r)rBrBXs eU #%   4 4UE BE_ O PP/?U kk!n-.! ^ !T!2T OF eV , 3 34E FF WWQ ( ("*<*<=xx#((8*<=>H zz(#H    *wwz3#;#;JHY#Z[* x%6QV W& 776: & - -e 4& 4 JNFBFF 3F4 -4 -s 9G""G,c|rXtj|tjs1tj|tjr d}t ||}nhtj|tj stj |}|}n,tjd|jj}|}t|tj|fS)Na jax.numpy.var does not yet support real dtype parameters when computing the variance of an array of complex values. The semantics of numpy.var seem unclear in this case. Please comment on https://github.com/jax-ml/jax/issues/2283 if this behavior is important to you.r) rrAr&rrMrkr!rrr:r>)rr:msgrvs r)rGrGs   eR%7%7 8'2#5#56!c sO   Wbjj 1%%g.ehhq'"''--e! & '% 88r+c ||}n t|tr|dk7r tdtd|}t |t |||||||S)a Compute the standard deviation along a given axis. JAX implementation of :func:`numpy.std`. Args: a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the standard deviation is computed. If None, standard deviaiton is computed along all the axes. dtype: The type of the output array. Default=None. ddof: int, default=0. Degrees of freedom. The divisor in the standard deviation computation is ``N-ddof``, ``N`` is number of elements along given axis. keepdims: bool, default=False. If true, reduced axes are left in the result with size 1. where: optional, boolean array, default=None. The elements to be used in the standard deviation. Array should be broadcast compatible to the input. mean: optional, mean of the input array, computed along the given axis. If provided, it will be used to compute the standard deviation instead of computing it from the input array. If specified, mean must be broadcast-compatible with the input array. In the general case, this can be achieved by computing the mean with ``keepdims=True`` and ``axis`` matching this function's ``axis`` argument. correction: int or float, default=None. Alternative name for ``ddof``. Both ddof and correction can't be provided simultaneously. out: Unused by JAX. Returns: An array of the standard deviation along the given axis. See also: - :func:`jax.numpy.var`: Compute the variance of array elements over given axis. - :func:`jax.numpy.mean`: Compute the mean of array elements over a given axis. - :func:`jax.numpy.nanvar`: Compute the variance along a given axis, ignoring NaNs values. - :func:`jax.numpy.nanstd`: Computed the standard deviation of a given axis, ignoring NaN values. Examples: By default, ``jnp.std`` computes the standard deviation along all axes. >>> x = jnp.array([[1, 3, 4, 2], ... [4, 2, 5, 3], ... [5, 4, 2, 3]]) >>> with jnp.printoptions(precision=2, suppress=True): ... jnp.std(x) Array(1.21, dtype=float32) If ``axis=0``, computes along axis 0. >>> with jnp.printoptions(precision=2, suppress=True): ... print(jnp.std(x, axis=0)) [1.7 0.82 1.25 0.47] To preserve the dimensions of input, you can set ``keepdims=True``. >>> with jnp.printoptions(precision=2, suppress=True): ... print(jnp.std(x, axis=0, keepdims=True)) [[1.7 0.82 1.25 0.47]] If ``ddof=1``: >>> with jnp.printoptions(precision=2, suppress=True): ... print(jnp.std(x, axis=0, keepdims=True, ddof=1)) [[2.08 1. 1.53 0.58]] To include specific elements of the array to compute standard deviation, you can use ``where``. >>> where = jnp.array([[1, 0, 1, 0], ... [0, 1, 0, 1], ... [1, 1, 1, 0]], dtype=bool) >>> jnp.std(x, axis=0, keepdims=True, where=where) Array([[2., 1., 1., 0.]], dtype=float32) rr>stdrUr:rVr<rWrOr)rrrMr_stdrrCs r)rPrPs^^J dC DAI L MMua ! a+D1CT^iq &&r+rQc td|}|Htj|d}tj|tj st d|| tdtjt|||||||S)NrPz/dtype argument to jnp.std must be inexact; got z/The 'out' argument to jnp.std is not supported.)rUr:r<rWrOr) rPrrjrAr&rkrMrfrsqrtr?)r2rUr:rVr<rWrOrs r)rRrRs eU #%   4 4UE BE   UBJJ / HP QQ_ O PP #ad%J'u4A BBr+cHtd|}t|t|||S)aReturn the peak-to-peak range along a given axis. JAX implementation of :func:`numpy.ptp`. Args: a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the range is computed. If None, the range is computed on the flattened array. keepdims: bool, default=False. If true, reduced axes are left in the result with size 1. out: Unused by JAX. Returns: An array with the range of elements along specified axis of input. Examples: By default, ``jnp.ptp`` computes the range along all axes. >>> x = jnp.array([[1, 3, 5, 2], ... [4, 6, 8, 1], ... [7, 9, 3, 4]]) >>> jnp.ptp(x) Array(8, dtype=int32) If ``axis=1``, computes the range along axis 1. >>> jnp.ptp(x, axis=1) Array([4, 7, 6], dtype=int32) To preserve the dimensions of input, you can set ``keepdims=True``. >>> jnp.ptp(x, axis=1, keepdims=True) Array([[4], [7], [6]], dtype=int32) ptp)r_ptpr)r2rUrVrWs r)rVrVs(Nua ! a&t,c8 <>> x = jnp.array([[1, 0, 0, 0], ... [0, 0, 1, 0], ... [1, 1, 1, 0]]) >>> jnp.count_nonzero(x) Array(5, dtype=int32) If ``axis=1``, counts along axis 1. >>> jnp.count_nonzero(x, axis=1) Array([1, 1, 3], dtype=int32) To preserve the dimensions of input, you can set ``keepdims=True``. >>> jnp.count_nonzero(x, axis=1, keepdims=True) Array([[1], [1], [3]], dtype=int32) count_nonzerorrE)rrrne_constrrC)r2rUrWs r)r[r[/sHN*! SVVAszz!Q' (t++- BBr+c t||}t||}tj|jt j s ||f|||d|S|ttj|t|||f|||d|} |rNtttj|||tj|t j| S| S)Nrr)rrPrrAr:r&rkrr_isnanrmrr]rK) r2rN jnp_reductionrrnan_if_all_nanrUrWrOkwargsrVs r)_nan_reductionrc[stQ! dE "%   177BJJ /  P P PPfSZZ],?8,LaP K(% KCI K# #cjjm$B**Q' .. Jr+c Tt|dttj|du||||| S)aJ Return the minimum of the array elements along a given axis, ignoring NaNs. JAX implementation of :func:`numpy.nanmin`. Args: a: Input array. axis: int or sequence of ints, default=None. Axis along which the minimum is computed. If None, the minimum is computed along the flattened array. keepdims: bool, default=False. If True, reduced axes are left in the result with size 1. initial: int or array, default=None. Initial value for the minimum. where: array of boolean dtype, default=None. The elements to be used in the minimum. Array should be broadcast compatible to the input. ``initial`` must be specified when ``where`` is used. out: Unused by JAX. Returns: An array of minimum values along the given axis, ignoring NaNs. If all values are NaNs along the given axis, returns ``nan``. See also: - :func:`jax.numpy.nanmax`: Compute the maximum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nansum`: Compute the sum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nanprod`: Compute the product of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nanmean`: Compute the mean of array elements along a given axis, ignoring NaNs. Examples: By default, ``jnp.nanmin`` computes the minimum of elements along the flattened array. >>> nan = jnp.nan >>> x = jnp.array([[1, nan, 4, 5], ... [nan, -2, nan, -4], ... [2, 1, 3, nan]]) >>> jnp.nanmin(x) Array(-4., dtype=float32) If ``axis=1``, the maximum will be computed along axis 1. >>> jnp.nanmin(x, axis=1) Array([ 1., -4., 1.], dtype=float32) If ``keepdims=True``, ``ndim`` of the output will be same of that of the input. >>> jnp.nanmin(x, axis=1, keepdims=True) Array([[ 1.], [-4.], [ 1.]], dtype=float32) To include only specific elements in computing the maximum, you can use ``where``. It can either have same dimension as input >>> where=jnp.array([[0, 0, 1, 0], ... [0, 0, 1, 1], ... [1, 1, 1, 0]], dtype=bool) >>> jnp.nanmin(x, axis=1, keepdims=True, initial=0, where=where) Array([[ 0.], [-4.], [ 0.]], dtype=float32) or must be broadcast compatible with input. >>> where = jnp.array([[False], ... [True], ... [False]]) >>> jnp.nanmin(x, axis=0, keepdims=True, initial=0, where=where) Array([[ 0., -2., 0., -4.]], dtype=float32) nanminNrarUrVrWrXrO)rcrr&rrs r)rerems0\ 8S"&&D!sX 'u 66r+c Vt|dttj |du||||| S)aC Return the maximum of the array elements along a given axis, ignoring NaNs. JAX implementation of :func:`numpy.nanmax`. Args: a: Input array. axis: int or sequence of ints, default=None. Axis along which the maximum is computed. If None, the maximum is computed along the flattened array. keepdims: bool, default=False. If True, reduced axes are left in the result with size 1. initial: int or array, default=None. Initial value for the maximum. where: array of boolean dtype, default=None. The elements to be used in the maximum. Array should be broadcast compatible to the input. ``initial`` must be specified when ``where`` is used. out: Unused by JAX. Returns: An array of maximum values along the given axis, ignoring NaNs. If all values are NaNs along the given axis, returns ``nan``. See also: - :func:`jax.numpy.nanmin`: Compute the minimum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nansum`: Compute the sum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nanprod`: Compute the product of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nanmean`: Compute the mean of array elements along a given axis, ignoring NaNs. Examples: By default, ``jnp.nanmax`` computes the maximum of elements along the flattened array. >>> nan = jnp.nan >>> x = jnp.array([[8, nan, 4, 6], ... [nan, -2, nan, -4], ... [-2, 1, 7, nan]]) >>> jnp.nanmax(x) Array(8., dtype=float32) If ``axis=1``, the maximum will be computed along axis 1. >>> jnp.nanmax(x, axis=1) Array([ 8., -2., 7.], dtype=float32) If ``keepdims=True``, ``ndim`` of the output will be same of that of the input. >>> jnp.nanmax(x, axis=1, keepdims=True) Array([[ 8.], [-2.], [ 7.]], dtype=float32) To include only specific elements in computing the maximum, you can use ``where``. It can either have same dimension as input >>> where=jnp.array([[0, 0, 1, 0], ... [0, 0, 1, 1], ... [1, 1, 1, 0]], dtype=bool) >>> jnp.nanmax(x, axis=1, keepdims=True, initial=0, where=where) Array([[4.], [0.], [7.]], dtype=float32) or must be broadcast compatible with input. >>> where = jnp.array([[True], ... [False], ... [False]]) >>> jnp.nanmax(x, axis=0, keepdims=True, initial=0, where=where) Array([[8., 0., 4., 6.]], dtype=float32) nanmaxNrf)rcrr&rrs r)rhrhs2\ 8S266''T/!sX 'u 66r+c f|tj|d}t|dtdd|||||| S)a Return the sum of the array elements along a given axis, ignoring NaNs. JAX implementation of :func:`numpy.nansum`. Args: a: Input array. axis: int or sequence of ints, default=None. Axis along which the sum is computed. If None, the sum is computed along the flattened array. dtype: The type of the output array. Default=None. keepdims: bool, default=False. If True, reduced axes are left in the result with size 1. initial: int or array, default=None. Initial value for the sum. where: array of boolean dtype, default=None. The elements to be used in the sum. Array should be broadcast compatible to the input. out: Unused by JAX. Returns: An array containing the sum of array elements along the given axis, ignoring NaNs. If all elements along the given axis are NaNs, returns 0. See also: - :func:`jax.numpy.nanmin`: Compute the minimum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nanmax`: Compute the maximum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nanprod`: Compute the product of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nanmean`: Compute the mean of array elements along a given axis, ignoring NaNs. Examples: By default, ``jnp.nansum`` computes the sum of elements along the flattened array. >>> nan = jnp.nan >>> x = jnp.array([[3, nan, 4, 5], ... [nan, -2, nan, 7], ... [2, 1, 6, nan]]) >>> jnp.nansum(x) Array(26., dtype=float32) If ``axis=1``, the sum will be computed along axis 1. >>> jnp.nansum(x, axis=1) Array([12., 5., 9.], dtype=float32) If ``keepdims=True``, ``ndim`` of the output will be same of that of the input. >>> jnp.nansum(x, axis=1, keepdims=True) Array([[12.], [ 5.], [ 9.]], dtype=float32) To include only specific elements in computing the sum, you can use ``where``. >>> where=jnp.array([[1, 0, 1, 0], ... [0, 0, 1, 1], ... [1, 1, 1, 0]], dtype=bool) >>> jnp.nansum(x, axis=1, keepdims=True, where=where) Array([[7.], [7.], [9.]], dtype=float32) If ``where`` is ``False`` at all elements, ``jnp.nansum`` returns 0 along the given axis. >>> where = jnp.array([[False], ... [False], ... [False]]) >>> jnp.nansum(x, axis=0, keepdims=True, where=where) Array([[0., 0., 0., 0.]], dtype=float32) nanprodnansumrFrarUr:rVrWrXrO)rrjrcrrs r)rkrksB^   4 4UI FE 8S!E!C( 'u 66r+c f|tj|d}t|dtdd|||||| S)aE Return the product of the array elements along a given axis, ignoring NaNs. JAX implementation of :func:`numpy.nanprod`. Args: a: Input array. axis: int or sequence of ints, default=None. Axis along which the product is computed. If None, the product is computed along the flattened array. dtype: The type of the output array. Default=None. keepdims: bool, default=False. If True, reduced axes are left in the result with size 1. initial: int or array, default=None. Initial value for the product. where: array of boolean dtype, default=None. The elements to be used in the product. Array should be broadcast compatible to the input. out: Unused by JAX. Returns: An array containing the product of array elements along the given axis, ignoring NaNs. If all elements along the given axis are NaNs, returns 1. See also: - :func:`jax.numpy.nanmin`: Compute the minimum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nanmax`: Compute the maximum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nansum`: Compute the sum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nanmean`: Compute the mean of array elements along a given axis, ignoring NaNs. Examples: By default, ``jnp.nanprod`` computes the product of elements along the flattened array. >>> nan = jnp.nan >>> x = jnp.array([[nan, 3, 4, nan], ... [5, nan, 1, 3], ... [2, 1, nan, 1]]) >>> jnp.nanprod(x) Array(360., dtype=float32) If ``axis=1``, the product will be computed along axis 1. >>> jnp.nanprod(x, axis=1) Array([12., 15., 2.], dtype=float32) If ``keepdims=True``, ``ndim`` of the output will be same of that of the input. >>> jnp.nanprod(x, axis=1, keepdims=True) Array([[12.], [15.], [ 2.]], dtype=float32) To include only specific elements in computing the maximum, you can use ``where``. >>> where=jnp.array([[1, 0, 1, 0], ... [0, 0, 1, 1], ... [1, 1, 1, 0]], dtype=bool) >>> jnp.nanprod(x, axis=1, keepdims=True, where=where) Array([[4.], [3.], [2.]], dtype=float32) If ``where`` is ``False`` at all elements, ``jnp.nanprod`` returns 1 along the given axis. >>> where = jnp.array([[False], ... [False], ... [False]]) >>> jnp.nanprod(x, axis=0, keepdims=True, where=where) Array([[1., 1., 1., 1.]], dtype=float32) rjr`Frl)rrjrcrrs r)rjrjisB\   4 4UI FE 9dAe!C( 'u 66r+c Rtd|}td|}| tdtj|j t js.tj|j t jrt||||||S| tj|j }ntj|d}tjtj|}t|||||}tj t#||||||}|S)a Return the mean of the array elements along a given axis, ignoring NaNs. JAX implementation of :func:`numpy.nanmean`. Args: a: Input array. axis: int or sequence of ints, default=None. Axis along which the mean is computed. If None, the mean is computed along the flattened array. dtype: The type of the output array. Default=None. keepdims: bool, default=False. If True, reduced axes are left in the result with size 1. where: array of boolean dtype, default=None. The elements to be used in computing mean. Array should be broadcast compatible to the input. out: Unused by JAX. Returns: An array containing the mean of array elements along the given axis, ignoring NaNs. If all elements along the given axis are NaNs, returns ``nan``. See also: - :func:`jax.numpy.nanmin`: Compute the minimum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nanmax`: Compute the maximum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nansum`: Compute the sum of array elements along a given axis, ignoring NaNs. - :func:`jax.numpy.nanprod`: Compute the product of array elements along a given axis, ignoring NaNs. Examples: By default, ``jnp.nanmean`` computes the mean of elements along the flattened array. >>> nan = jnp.nan >>> x = jnp.array([[2, nan, 4, 3], ... [nan, -2, nan, 9], ... [4, -7, 6, nan]]) >>> jnp.nanmean(x) Array(2.375, dtype=float32) If ``axis=1``, mean will be computed along axis 1. >>> jnp.nanmean(x, axis=1) Array([3. , 3.5, 1. ], dtype=float32) If ``keepdims=True``, ``ndim`` of the output will be same of that of the input. >>> jnp.nanmean(x, axis=1, keepdims=True) Array([[3. ], [3.5], [1. ]], dtype=float32) ``where`` can be used to include only specific elements in computing the mean. >>> where = jnp.array([[1, 0, 1, 0], ... [0, 0, 1, 1], ... [1, 1, 0, 1]], dtype=bool) >>> jnp.nanmean(x, axis=1, keepdims=True, where=where) Array([[ 3. ], [ 9. ], [-1.5]], dtype=float32) If ``where`` is ``False`` at all elements, ``jnp.nanmean`` returns ``nan`` along the given axis. >>> where = jnp.array([[False], ... [False], ... [False]]) >>> jnp.nanmean(x, axis=0, keepdims=True, where=where) Array([[nan, nan, nan, nan]], dtype=float32) nanmeanz3The 'out' argument to jnp.nanmean is not supported.rr)rUr:rWrOr )rrPrfrrAr:r&rBrHrr!rjr bitwise_notr_rr"rk) r2rUr:rVrWrOnan_maskr#tds r)rorosXy!$! i '%_ S TT qww)V->->qww -S 4XU ;; ]  # #AGG ,E  4 4UF CE __SZZ] +(8$eheT* wwvaUXUKZX" )r+c td|}td|}|tj|d}| t dt ||||||||S)a_ Compute the variance of array elements along a given axis, ignoring NaNs. JAX implementation of :func:`numpy.nanvar`. Args: a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the variance is computed. If None, variance is computed along flattened array. dtype: The type of the output array. Default=None. ddof: int, default=0. Degrees of freedom. The divisor in the variance computation is ``N-ddof``, ``N`` is number of elements along given axis. keepdims: bool, default=False. If true, reduced axes are left in the result with size 1. where: optional, boolean array, default=None. The elements to be used in the variance. Array should be broadcast compatible to the input. mean: optional, mean of the input array, computed along the given axis. If provided, it will be used to compute the variance instead of computing it from the input array. If specified, mean must be broadcast-compatible with the input array. In the general case, this can be achieved by computing the mean with ``keepdims=True`` and ``axis`` matching this function's ``axis`` argument. out: Unused by JAX. Returns: An array containing the variance of array elements along specified axis. If all elements along the given axis are NaNs, returns ``nan``. See also: - :func:`jax.numpy.nanmean`: Compute the mean of array elements over a given axis, ignoring NaNs. - :func:`jax.numpy.nanstd`: Computed the standard deviation of a given axis, ignoring NaNs. - :func:`jax.numpy.var`: Compute the variance of array elements along a given axis. Examples: By default, ``jnp.nanvar`` computes the variance along all axes. >>> nan = jnp.nan >>> x = jnp.array([[1, nan, 4, 3], ... [nan, 2, nan, 9], ... [4, 8, 6, nan]]) >>> jnp.nanvar(x) Array(6.984375, dtype=float32) If ``axis=1``, variance is computed along axis 1. >>> with jnp.printoptions(precision=2, suppress=True): ... print(jnp.nanvar(x, axis=1)) [ 1.56 12.25 2.67] To preserve the dimensions of input, you can set ``keepdims=True``. >>> with jnp.printoptions(precision=2, suppress=True): ... print(jnp.nanvar(x, axis=1, keepdims=True)) [[ 1.56] [12.25] [ 2.67]] If ``ddof=1``: >>> with jnp.printoptions(precision=2, suppress=True): ... print(jnp.nanvar(x, axis=1, keepdims=True, ddof=1)) [[ 2.33] [24.5 ] [ 4. ]] To include specific elements of the array to compute variance, you can use ``where``. >>> where = jnp.array([[1, 0, 1, 0], ... [0, 1, 1, 0], ... [1, 1, 0, 1]], dtype=bool) >>> jnp.nanvar(x, axis=1, keepdims=True, where=where) Array([[2.25], [0. ], [4. ]], dtype=float32) nanvarz2The 'out' argument to jnp.nanvar is not supported.rUr:rVrDrWrOrA)rrPrrjrf_nanvarr2rUr:rVrDrWrOrs r)rtrtsadx#! h &%   4 4UH EE_ R SS U$Y^gk llr+ruct|j|\}}tj|j |}|t |||d|}nt d|j |}ttj|dtj||} tj| jtjr=tjtj| tj | } ntj"| } t%tj&tj||||} | |z } tj(| tj*| } t%| |||} t| tj,| } t| d| } tj.| tj0| | j} tj0| |S)NTr rtrr)rWrOr`)rGr:rrorrorrr_rrrAr&rrrrHrIrrple_zerorKr"rl)r2rUr:rVrDrWrOrArvrLr#normalizer_maskrwdivisors r)rvrvtsy0?U kk!n-.! ^ Q$5E RF h / 6 67H IF CJJqM1cgga&8 9( x~~r'9'9:xx#((8*<=>Hzz(#H3??3::a=1xu>*D *FF:syy'<=/ x >& /2666 2& ?Az 2' 776333GV\\J K& ! !&% 00r+c td|}td|}|tj|d}| t dt j t|||||||S)a Compute the standard deviation along a given axis, ignoring NaNs. JAX implementation of :func:`numpy.nanstd`. Args: a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the standard deviation is computed. If None, standard deviaiton is computed along flattened array. dtype: The type of the output array. Default=None. ddof: int, default=0. Degrees of freedom. The divisor in the standard deviation computation is ``N-ddof``, ``N`` is number of elements along given axis. keepdims: bool, default=False. If true, reduced axes are left in the result with size 1. where: optional, boolean array, default=None. The elements to be used in the standard deviation. Array should be broadcast compatible to the input. mean: optional, mean of the input array, computed along the given axis. If provided, it will be used to compute the standard deviation instead of computing it from the input array. If specified, mean must be broadcast-compatible with the input array. In the general case, this can be achieved by computing the mean with ``keepdims=True`` and ``axis`` matching this function's ``axis`` argument. out: Unused by JAX. Returns: An array containing the standard deviation of array elements along the given axis. If all elements along the given axis are NaNs, returns ``nan``. See also: - :func:`jax.numpy.nanmean`: Compute the mean of array elements over a given axis, ignoring NaNs. - :func:`jax.numpy.nanvar`: Compute the variance along the given axis, ignoring NaNs values. - :func:`jax.numpy.std`: Computed the standard deviation along the given axis. Examples: By default, ``jnp.nanstd`` computes the standard deviation along flattened array. >>> nan = jnp.nan >>> x = jnp.array([[3, nan, 4, 5], ... [nan, 2, nan, 7], ... [2, 1, 6, nan]]) >>> jnp.nanstd(x) Array(1.9843135, dtype=float32) If ``axis=0``, computes standard deviation along axis 0. >>> jnp.nanstd(x, axis=0) Array([0.5, 0.5, 1. , 1. ], dtype=float32) To preserve the dimensions of input, you can set ``keepdims=True``. >>> jnp.nanstd(x, axis=0, keepdims=True) Array([[0.5, 0.5, 1. , 1. ]], dtype=float32) If ``ddof=1``: >>> with jnp.printoptions(precision=2, suppress=True): ... print(jnp.nanstd(x, axis=0, keepdims=True, ddof=1)) [[0.71 0.71 1.41 1.41]] To include specific elements of the array to compute standard deviation, you can use ``where``. >>> where=jnp.array([[1, 0, 1, 0], ... [0, 1, 0, 1], ... [1, 1, 0, 1]], dtype=bool) >>> jnp.nanstd(x, axis=0, keepdims=True, where=where) Array([[0.5, 0.5, 0. , 0. ]], dtype=float32) nanstdz2The 'out' argument to jnp.nanstd is not supported.)rUr:rDrWrOr)rrPrrjrfrrTrtrws r)r~r~spTx#! h &%   4 4UH EE_ R SS &U"*%dD EEr+c$eZdZ d ddZy)CumulativeReductionNcyr@ra)selfr2rUr:rVs r)__call__zCumulativeReduction.__call__sKNr+NNN) r2rrUAxisr:DTypeLike | NonerVNonereturnr)__name__ __module__ __qualname__rrar+r)rrs)04;?O&O48ODIOr+rc t||}|td|d| t|r*tj|t j |f}|d}tt j|} t| } t|| }|r4ttj|tj|||}|j} |4| } |s$tj | t j"rQt%| } nEtj&||} tj | t j"r t%| } | t j"k7rE|t j"k(r2tj(|j+t j"}tj,|| }|||} |Htj |t j"r$tj,| t j"} | S)z7Helper function for implementing cumulative reductions.r]z is not supportedr)rrfr*rr7r&rrrdrrrr_r]r:rrArBrIrjrorrl)rN reductionr2rUr:rVfill_nan fill_valuer[rnum_dimsa_type result_typerws r)_cumulative_reductionrs tQ!_  :4&@QR SS \Yq\ A }%A \ D ! ' \( 4 *$ szz!}cjjJ7;A''& ]K6,,["((C*;7k::5$GK bhh/*;7k rxxERXX- Abhh'A q+.! Q & 6,,UBHH=  % %fbhh 7F -r+rUr:c>tdtj||||S)aYCumulative sum of elements along an axis. JAX implementation of :func:`numpy.cumsum`. Args: a: N-dimensional array to be accumulated. axis: integer axis along which to accumulate. If None (default), then array will be flattened and accumulated along the flattened axis. dtype: optionally specify the dtype of the output. If not specified, then the output dtype will match the input dtype. out: unused by JAX Returns: An array containing the accumulated sum along the given axis. See also: - :func:`jax.numpy.cumulative_sum`: cumulative sum via the array API standard. - :meth:`jax.numpy.add.accumulate`: cumulative sum via ufunc methods. - :func:`jax.numpy.nancumsum`: cumulative sum ignoring NaN values. - :func:`jax.numpy.sum`: sum along axis Examples: >>> x = jnp.array([[1, 2, 3], ... [4, 5, 6]]) >>> jnp.cumsum(x) # flattened cumulative sum Array([ 1, 3, 6, 10, 15, 21], dtype=int32) >>> jnp.cumsum(x, axis=1) # cumulative sum along axis 1 Array([[ 1, 3, 6], [ 4, 9, 15]], dtype=int32) cumsumrrrr2rUr:rVs r)rrs D x)<)tdtj||||S)a8Cumulative product of elements along an axis. JAX implementation of :func:`numpy.cumprod`. Args: a: N-dimensional array to be accumulated. axis: integer axis along which to accumulate. If None (default), then array will be flattened and accumulated along the flattened axis. dtype: optionally specify the dtype of the output. If not specified, then the output dtype will match the input dtype. out: unused by JAX Returns: An array containing the accumulated product along the given axis. See also: - :meth:`jax.numpy.multiply.accumulate`: cumulative product via ufunc methods. - :func:`jax.numpy.nancumprod`: cumulative product ignoring NaN values. - :func:`jax.numpy.prod`: product along axis Examples: >>> x = jnp.array([[1, 2, 3], ... [4, 5, 6]]) >>> jnp.cumprod(x) # flattened cumulative product Array([ 1, 2, 6, 24, 120, 720], dtype=int32) >>> jnp.cumprod(x, axis=1) # cumulative product along axis 1 Array([[ 1, 2, 6], [ 4, 20, 120]], dtype=int32) cumprodrrrrs r)rr:s!B y,*>*>4PS TTr+c Dtdtj||||ddS)avCumulative sum of elements along an axis, ignoring NaN values. JAX implementation of :func:`numpy.nancumsum`. Args: a: N-dimensional array to be accumulated. axis: integer axis along which to accumulate. If None (default), then array will be flattened and accumulated along the flattened axis. dtype: optionally specify the dtype of the output. If not specified, then the output dtype will match the input dtype. out: unused by JAX Returns: An array containing the accumulated sum along the given axis. See also: - :func:`jax.numpy.cumsum`: cumulative sum without ignoring NaN values. - :func:`jax.numpy.cumulative_sum`: cumulative sum via the array API standard. - :meth:`jax.numpy.add.accumulate`: cumulative sum via ufunc methods. - :func:`jax.numpy.sum`: sum along axis Examples: >>> x = jnp.array([[1., 2., jnp.nan], ... [4., jnp.nan, 6.]]) The standard cumulative sum will propagate NaN values: >>> jnp.cumsum(x) Array([ 1., 3., nan, nan, nan, nan], dtype=float32) :func:`~jax.numpy.nancumsum` will ignore NaN values, effectively replacing them with zeros: >>> jnp.nancumsum(x) Array([ 1., 3., 3., 7., 7., 13.], dtype=float32) Cumulative sum along axis 1: >>> jnp.nancumsum(x, axis=1) Array([[ 1., 3., 3.], [ 4., 4., 10.]], dtype=float32) nancumsumTrrrrrs r)rr^s*\ {L,?,?D%QT(, <>> x = jnp.array([[1., 2., jnp.nan], ... [4., jnp.nan, 6.]]) The standard cumulative product will propagate NaN values: >>> jnp.cumprod(x) Array([ 1., 2., nan, nan, nan, nan], dtype=float32) :func:`~jax.numpy.nancumprod` will ignore NaN values, effectively replacing them with ones: >>> jnp.nancumprod(x) Array([ 1., 2., 2., 8., 8., 48.], dtype=float32) Cumulative product along axis 1: >>> jnp.nancumprod(x, axis=1) Array([[ 1., 2., 2.], [ 4., 4., 24.]], dtype=float32) nancumprodTr`rrrs r)rrs*Z |\-A-A1dESV(, <>> x = jnp.array([[1, 2, 3], ... [4, 5, 6]]) >>> jnp.cumulative_sum(x, axis=1) Array([[ 1, 3, 6], [ 4, 9, 15]], dtype=int32) >>> jnp.cumulative_sum(x, axis=1, include_initial=True) Array([[ 0, 1, 3, 6], [ 0, 4, 9, 15]], dtype=int32) cumulative_sumrzhThe input must be non-scalar to take a cumulative sum, however a scalar value or scalar array was given.r`The input array has rank o, however axis was not set to an explicit value. The axis argument is only optional for one-dimensional arrays.rr= dimension) rr.rMrrrjrrrdr concatenater4r:rrUr:rrV zeros_shapes r)rrsF'+!VVq[ 0  \ Dvvz  #AFF8,   4 ($   4 4U ;Eqt59#qww-KK // xx Qcii0#6 C *r+ctd|}|jdk(r td|*d}|jdkDrtd|jdt||j}|t j |}t dtj|||}|rSt|j}d||<tjtj|d|j|g|}|S) a.Cumulative product along the axis of an array. JAX implementation of :func:`numpy.cumulative_prod`. Args: x: N-dimensional array axis: integer axis along which to accumulate. If ``x`` is one-dimensional, this argument is optional and defaults to zero. dtype: optional dtype of the output. include_initial: if True, then include the initial value in the cumulative product. Default is False. Returns: An array containing the accumulated values. See Also: - :func:`jax.numpy.cumprod`: alternative API for cumulative product. - :func:`jax.numpy.nancumprod`: cumulative product while ignoring NaN values. - :func:`jax.numpy.multiply.accumulate`: cumulative product via the ufunc API. Examples: >>> x = jnp.array([[1, 2, 3], ... [4, 5, 6]]) >>> jnp.cumulative_prod(x, axis=1) Array([[ 1, 2, 6], [ 4, 20, 120]], dtype=int32) >>> jnp.cumulative_prod(x, axis=1, include_initial=True) Array([[ 1, 1, 2, 6], [ 1, 4, 20, 120]], dtype=int32) cumulative_prodrzlThe input must be non-scalar to take a cumulative product, however a scalar value or scalar array was given.r`rrr=r)rr.rMrrrjrrrrrdrrr4r:rs r)rr sF(!,!VVq[ 0  \ Dvvz  #AFF8,   4 ($   4 4U ;E/1E1Eq$PUV#qww-KK // xx Qcii0#6 C *r+)rUoverwrite_input interpolationrWmethod)rctd||\}}|s| tdt|ts t dt t j|t j||||dS)aCompute the quantile of the data along the specified axis. JAX implementation of :func:`numpy.quantile`. Args: a: N-dimensional array input. q: scalar or 1-dimensional array specifying the desired quantiles. ``q`` should contain floating-point values between ``0.0`` and ``1.0``. axis: optional axis or tuple of axes along which to compute the quantile out: not implemented by JAX; will error if not None overwrite_input: not implemented by JAX; will error if not False method: specify the interpolation method to use. Options are one of ``["linear", "lower", "higher", "midpoint", "nearest"]``. default is ``linear``. keepdims: if True, then the returned array will have the same number of dimensions as the input. Default is False. Returns: An array containing the specified quantiles along the specified axes. See also: - :func:`jax.numpy.nanquantile`: compute the quantile while ignoring NaNs - :func:`jax.numpy.percentile`: compute the percentile (0-100) Examples: Computing the median and quartiles of an array, with linear interpolation: >>> x = jnp.arange(10) >>> q = jnp.array([0.25, 0.5, 0.75]) >>> jnp.quantile(x, q) Array([2.25, 4.5 , 6.75], dtype=float32) Computing the quartiles using nearest-value interpolation: >>> jnp.quantile(x, q, method='nearest') Array([2., 4., 7.], dtype=float32) quantilezGjax.numpy.quantile does not support overwrite_input=True or out != NonezPquantile() argument interpolation was removed in JAX v0.8.0. Use method instead.FrrMrrr _quantilerror2qrUrVrrrWrs r)rrG ssT *a +$!Q & '' M= 1 3 44 3;;q>3;;q>45 QQr+ctd||\}}|s| d}t|t|ts t dt t j|t j||||dS)aCompute the quantile of the data along the specified axis, ignoring NaNs. JAX implementation of :func:`numpy.nanquantile`. Args: a: N-dimensional array input. q: scalar or 1-dimensional array specifying the desired quantiles. ``q`` should contain floating-point values between ``0.0`` and ``1.0``. axis: optional axis or tuple of axes along which to compute the quantile out: not implemented by JAX; will error if not None overwrite_input: not implemented by JAX; will error if not False method: specify the interpolation method to use. Options are one of ``["linear", "lower", "higher", "midpoint", "nearest"]``. default is ``linear``. keepdims: if True, then the returned array will have the same number of dimensions as the input. Default is False. Returns: An array containing the specified quantiles along the specified axes. See also: - :func:`jax.numpy.quantile`: compute the quantile without ignoring nans - :func:`jax.numpy.nanpercentile`: compute the percentile (0-100) Examples: Computing the median and quartiles of a 1D array: >>> x = jnp.array([0, 1, 2, jnp.nan, 3, 4, 5, 6]) >>> q = jnp.array([0.25, 0.5, 0.75]) Because of the NaN value, :func:`jax.numpy.quantile` returns all NaNs, while :func:`~jax.numpy.nanquantile` ignores them: >>> jnp.quantile(x, q) Array([nan, nan, nan], dtype=float32) >>> jnp.nanquantile(x, q) Array([1.5, 3. , 4.5], dtype=float32) nanquantilezJjax.numpy.nanquantile does not support overwrite_input=True or out != NonezSnanquantile() argument interpolation was removed in JAX v0.8.0. Use method instead.Tr) r2rrUrVrrrWrrNs r)rr| spV -A .$!Q C S/ M= 1 3 44 3;;q>3;;q>44 PPr+c Z!|dvr tdt|\}g}tj|jt j r td&|rdg|jz}|j}dnttr`t|j}|j!t!fdDtttk7r tdD]}d||< tt!tz }tt!} t!t#|D]\} } | | | | c| | <| | <tfdt!|jD} tfd t!|jD} t%j&|| t)j*| fz| }t-d |jnt-|j|j}|j}|dkDrtd |j|j}|rt/t%j0|t j2|}t%j4| }t7t%j8t%j0||j| }|j}t%j:|tt|t||z}t%j:|tt|}t%j<|t%j>|t%j@|d}t%jB|}t%jD|}t%j>||}t%j>t%j@|d|}t%jFt%j@|dt%jH||dz }t%jFt%j@|dt%jH||dz }t%jJ|tL}t%jJ|tL}||z}tt|Dcgc] }t%jNtL|||z"}}|r||<n|jQ||t|}||<|t|}ntSjTd5t/tWt%j0|dt j2|}dddt%j4| }t%jJ|t%jX|}t%j<||dz }t%jB|}t%jD|}t%j>||}t%j>t%j@|d|}t%jZt%j@|d||dz }t%jZt%j@|d||dz }t%jJ|tL}t%jJ|tL}t|}d|<t]j^tt||rt||znt||zdz |rdnff}t]j`||d||}t]j`||d||}|dk(rDt%jb||jd}t%jb||jd}|dk(rqt%jdt%j<|jg|j|t%j<|jg|j|}n|dk(r|}n|dk(r|}n|dk(rBt%jh|t%j@|d} t%jj| ||}nS|dk(r?t%j<t%jd||t%j@|d}ntd|d|r3|r1|dkDrt j|dg|}|j'|}t%jJ||jScc}w#1swYxYw)N)linearlowerhighermidpointnearestzHmethod can only be 'linear', 'lower', 'higher', 'midpoint', or 'nearest'zLquantile does not support complex input, as the operation is poorly defined.r`rc36K|]}t|ywr@r|)rbr0nds r)rez_quantile.. s:r"2r*:sz repeated axisc32K|]\}}|vs |ywr@rarbidxrrUs r)rez_quantile.. sTUSCtOqT c32K|]\}}|vs |ywr@rars r)rez_quantile.. sIec!SD[Irrz$q must be have rank <= 1, got shape rrEFTrra) offset_dimscollapsed_slice_dimsstart_index_map).N)dimension_numbers slice_sizes)r)broadcast_dimensionsrrrr?rzmethod=z not recognized)6rMrrrAr:r&rr.ravelrrrrdrrr/r6sortedrr7r5rrrr_rKsortrrprprrr]floorceilrrrlrbroadcasted_iotar0r rJr_dtypeclamp lax_slicingGatherDimensionNumbersgatherbroadcast_in_dimrrryr)"r2rrUrrW squash_nanskeepdimr0keepr3r5sdo_not_touch_shape touch_shapeq_shapeq_ndimrcountsshape_after_reductionlowhigh high_weight low_weight out_shaper2r low_value high_valuenrdnumsrwpredrs" ` @r)rr s= GG _ ``a "! ' qww 2 23 c dd \aff g  A D$177mG B :T: :D 3t9~T"  ''gbk uRy>CI %DeBiJ&,'B1%/]JqM"jmZ]BTi.@TTIy'9IIK A)TYY{-C,EEzRA R (D T166 *D GG' 66& aZ ;AGG9E FF GG'szz!}bffa(A d#A A/d!''T\ ]F"LL uVS!67&@A B DA __VU5=%9 :F 37763::a#345A ))A,C 88A;D''!S/KK3[AJ ''#**S!$cggc6A:&> ?C 773::dA&fqj(A BD " "3 ,C  # #D# .D//Ic"789 ; ! !#y#, ? ;E ;eDk ll4%,IE$K5<J   5 !J SZZ]=rvvq IaJ d#A    1 >A 1q5A ))A,C 88A;D''!S/KK3[AJ ))CJJsA&QU 3C 99SZZa($A 6D " "3 ,C  # #D# .Dw-KK  . .!)G vs7|f/Dq/HJK"*2wg  E ""1c)n/:>> x = jnp.array([0, 1, 2, 3, 4, 5, 6]) >>> q = jnp.array([25, 50, 75]) >>> jnp.percentile(x, q) Array([1.5, 3. , 4.5], dtype=float32) Computing the same percentiles with nearest rather than linear interpolation: >>> jnp.percentile(x, q, method='nearest') Array([1., 3., 4.], dtype=float32) percentilezRpercentile() argument interpolation was removed in JAX v0.8.0. Use method instead.drUrVrrrW)rrrrrrrs r)rr/ s_V ,1 -$!Qa "! M= 1 3 44 !QW4S/( 44r+c td||\}}t|\}|dz }t|ts t dt |||||||S)a Compute the percentile of the data along the specified axis, ignoring NaN values. JAX implementation of :func:`numpy.nanpercentile`. Args: a: N-dimensional array input. q: scalar or 1-dimensional array specifying the desired quantiles. ``q`` should contain integer or floating point values between ``0`` and ``100``. axis: optional axis or tuple of axes along which to compute the quantile out: not implemented by JAX; will error if not None overwrite_input: not implemented by JAX; will error if not False method: specify the interpolation method to use. Options are one of ``["linear", "lower", "higher", "midpoint", "nearest"]``. default is ``linear``. keepdims: if True, then the returned array will have the same number of dimensions as the input. Default is False. Returns: An array containing the specified percentiles along the specified axes. See also: - :func:`jax.numpy.nanquantile`: compute the nan-aware quantile (0.0-1.0) - :func:`jax.numpy.percentile`: compute the percentile without special handling of NaNs. Examples: Computing the median and quartiles of a 1D array: >>> x = jnp.array([0, 1, 2, jnp.nan, 3, 4, 5, 6]) >>> q = jnp.array([25, 50, 75]) Because of the NaN value, :func:`jax.numpy.percentile` returns all NaNs, while :func:`~jax.numpy.nanpercentile` ignores them: >>> jnp.percentile(x, q) Array([nan, nan, nan], dtype=float32) >>> jnp.nanpercentile(x, q) Array([1.5, 3. , 4.5], dtype=float32) nanpercentilerzUnanpercentile() argument interpolation was removed in JAX v0.8.0. Use method instead.r)rrrrrrrs r)rre sdZ /1a 0$!Qa "!#g! M= 1 3 44 Q#"X 77r+)rUrrWc >td|}t|d||||dS)aFReturn the median of array elements along a given axis. JAX implementation of :func:`numpy.median`. Args: a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the median to be computed. If None, median is computed for the flattened array. keepdims: bool, default=False. If true, reduced axes are left in the result with size 1. out: Unused by JAX. overwrite_input: Unused by JAX. Returns: An array of the median along the given axis. See also: - :func:`jax.numpy.mean`: Compute the mean of array elements over a given axis. - :func:`jax.numpy.max`: Compute the maximum of array elements over given axis. - :func:`jax.numpy.min`: Compute the minimum of array elements over given axis. Examples: By default, the median is computed for the flattened array. >>> x = jnp.array([[2, 4, 7, 1], ... [3, 5, 9, 2], ... [6, 1, 8, 3]]) >>> jnp.median(x) Array(3.5, dtype=float32) If ``axis=1``, the median is computed along axis 1. >>> jnp.median(x, axis=1) Array([3. , 4. , 4.5], dtype=float32) If ``keepdims=True``, ``ndim`` of the output is equal to that of the input. >>> jnp.median(x, axis=1, keepdims=True) Array([[3. ], [4. ], [4.5]], dtype=float32) medianrrrUrVrrWr)rrr2rUrVrrWs r)rr s-^x#! !Sto#J 88r+c >td|}t|d||||dS)a!Return the median of array elements along a given axis, ignoring NaNs. JAX implementation of :func:`numpy.nanmedian`. Args: a: input array. axis: optional, int or sequence of ints, default=None. Axis along which the median to be computed. If None, median is computed for the flattened array. keepdims: bool, default=False. If true, reduced axes are left in the result with size 1. out: Unused by JAX. overwrite_input: Unused by JAX. Returns: An array containing the median along the given axis, ignoring NaNs. If all elements along the given axis are NaNs, returns ``nan``. See also: - :func:`jax.numpy.nanmean`: Compute the mean of array elements over a given axis, ignoring NaNs. - :func:`jax.numpy.nanmax`: Compute the maximum of array elements over given axis, ignoring NaNs. - :func:`jax.numpy.nanmin`: Compute the minimum of array elements over given axis, ignoring NaNs. Examples: By default, the median is computed for the flattened array. >>> nan = jnp.nan >>> x = jnp.array([[2, nan, 7, nan], ... [nan, 5, 9, 2], ... [6, 1, nan, 3]]) >>> jnp.nanmedian(x) Array(4., dtype=float32) If ``axis=1``, the median is computed along axis 1. >>> jnp.nanmedian(x, axis=1) Array([4.5, 5. , 3. ], dtype=float32) If ``keepdims=True``, ``ndim`` of the output is equal to that of the input. >>> jnp.nanmedian(x, axis=1, keepdims=True) Array([[4.5], [5. ], [3. ]], dtype=float32) nanmedianrrr)rrrs r)rr s.h{A&! Q$C%4x& ((r+)r(rrrL)r2rr3rr4rrr)r:rrr)r:rrr)rNrrOArrayLike | Nonerz Array | None)"r2rrNrrq ReductionOprrrrQrLrRzCallable[[Array], Array] | NonerSzReductionOp | NonerTrLrUrr:rrVrrWrLrXrrYrrZzCallable[..., Array] | Noner[rLrr)r2rrUr)r2rrrrrz np.ndarray)rrrr)rrrr)rrrr)NNNFNNT)r2rrUrr:rrVrrWrLrXrrOrr[rLrr)NNNFNN)r2rrUrr:rrVrrWrLrXrrOrrr)NNFNN)r2rrUrrVrrWrLrXrrOrrr)NNF) r2rrUrrVrrWrLrOrrr)r2rrUzint | Sequence[int])NNNF)r2rrUrr:rrVrrWrLrOrrr) r2rrUrrWrLrOrr:r)r2rrUrr:rrVrrWrLrTrLrOrrr)NNFF) r2rrUrr&rr'zLiteral[False]rWrLrr)NN) r2rrUrr&rr'z Literal[True]rWrLrr) r2rrUrr&rr'rLrWrLrzArray | tuple[Array, Array])NNNrF)r2rrUrr:rrVrrDrrWrLrOrrrr<zint | float | Nonerr)r2rrUrr:rrVrr< int | floatrWrLrOrrArrr)rrr:rrztuple[DType, DType])r2rrUrr:rrVrr<rrWrLrOrrrrr) r2rrUrrVrrWrLrr) r2rrUrrVrrWrLrr)NF)r2rrUrrWrLrr)NFN)r2rrNrr`Callable[..., Array]rrrrarLrUrrWrLrOrrr)NNNFN)NNNrFNN)r2rrUrr:rrVrrDrrWrLrOrrrrr)r2rrUrr:rrVrrDrrWrLrOrrArrr)NFrF)rNrrrr2rrU int | Noner:rrVrrrLrrr[rLrrr) r2rrUrr:rrVrrr) rrrUrr:rrrLrr)NNFrF)r2rrrrUint | tuple[int, ...] | NonerVrrrLrrrWrLrrrr)r2rrrrUrrrrWrLrrLrr) r2rrUrrVrrrLrWrLrr) __future__rbuiltinscollections.abcrrr5rtypingrrrr r numpyr&jax._srcr r r rjax._src.numpy.utilrrrrr jax._src.laxrrrrrrrrjax._src.typingrrrrr jax._src.utilrrr r!exportrrirrr*r7r>rIrPrryrrhrmrrrrjitrrrrrrrrrrrrrrrrrrr r rrrrrr(r,r?rBrGrPrRrVrWr[rcrerhrkrjrortrvr~rrrrrrrrrrrrrrrrrar+r)r s#. ::<<&#+1/MMbb K || S(3- %&B   Sz3' (,:>-127 Dd %4*.>B(-AA $AA7AA+AA,0 AA  AA *: AA HL AA  AA 1A AA(AA!<AA"&AA38AAFV "G"5 , = JSWXKO38LP)-77,07)79I7#'7387Y7CGNRAEG8 G8$(G8;KG8G8:>G8JOG8G8V JSWXLP49MQ*.VV-1V*V:JV$(V49VYVDH,1EI"&H9H9%)H9"H92BH9 H9,1H9H9V -d;KO38LPVV,0V)V9IVUZV<V59<@"&GFGF)9GFGF+0GFGFR 6tDKO38LPVV,0V)V9IVUZVEV59<@"&FFFF)9FFFF+0FFFFP -d;=A!&IEIII2BINSI<I 59/5=A/5/5*:/5FK/5/5b -d;=A!&IEIII2BINSI<I 59/5=A/5/5*:/5FK/5/5d 6tDSW;@TX3!3483!13AQ3]b3E3 6tDRV;@TX3!3483!13AQ3]b3E3 6tDSW;@TX3!3483!13AQ3]b3E3 6tDSW;@TX3!3483!13AQ3]b3E3 6tDRV:?SW3 3373 03@P3\a3E3 6tDSW;@TX3!3483!13AQ3]b3E3KO27KOMM+/M(M8HMTYM.LP38LP88,08)89I8UZ8"6:<@"&++)9+++0++6:<@"&++)9+++0++ DH,1@H#'@H@H%)@H @H,1@H@HD     " TEI-2-1$(&*&*".3@ IM?DT$T8<TINT T K6;K#K/3K@EK K IM5:``.2`?Z` `IM5:6O6O.26O?Z6O6Op 9$GJN6;! ! /3! @[! H! HCG:?X("&)-X( X( #X(37X(X(.>X('X(38X(X(t 67#'4Q#'D$$'2$BF$ $1A$MR$8$N9*CG:?T&"&)-T& T& #T&37T&T&.>T&'T&38T&T&l 67#'4Q#'$ B B'2 BBF B  B/? BKP B8 B59'='=#('='=R -.26$)/-.-1#('B 'B-2'B/'BXY]&8<04EU!&$-.8BW W +;W GLW 8W t67X\+0DHVmVm$(Vm"Vm1AVmMRVm8Vmp'+dX\5%)d11$(1"13C1OT1667X\+0DHOEOE$(OE"OE1AOEMROE8OEdO(OJN45" * *.**"*+;*BF**)2* *(- *Z*+,07; S" S04 S@E S, SF*+-18<U#U15UAFU,UD*+/3:>-<%-<37-B5:!/8/8.2/8/8&+/8B/8d@AAE8=$5(5(155(5().5(B5(r+