o 6a^@szddlZddlZddlZddlZddlmZmZmZm Z m Z m Z m Z m Z mZmZmZmZmZmZddlmZmZmZmZmZejdkZddZGdddZGd d d ZGd d d ZGd ddZ GdddZ!GdddZ"GdddZ#GdddZ$GdddZ%GdddZ&GdddZ'GdddZ(Gdd d Z)Gd!d"d"Z*Gd#d$d$Z+Gd%d&d&Z,d'd(Z-dS))N)apply_along_axisapply_over_axes array_splitsplithsplitdsplitvsplitdstack column_stackkrontile expand_dimstake_along_axisput_along_axis)assert_ assert_equalassert_array_equal assert_raises assert_warnscstfdd}|S)z: hack in keepdims behavior into a function taking an axis cs.|fd|i|}|durd}tj||dS)Naxisrr)npr )arkwargsresfuncA/usr/lib/python3/dist-packages/numpy/lib/tests/test_shape_base.pywrappedsz_add_keepdims..wrapped) functoolswraps)rr rrr _add_keepdimssr#c@,eZdZddZddZddZddZd S) TestTakeAlongAxisc Csddlm}|ddd}tjtjtfttjttjtfttj ttj tftj tj tddfg}|D]0\}}}t t|jdgD] }||fd |i|}||fd |i|} t|t|| |d qEq6dS) z- Test it translates from arg to rrand)kthNrr) numpy.randomr'rsortargsortdictr#minargminmaxargmax partition argpartitionlistrangendimrr) selfr'rfuncsrargfuncrra_funcai_funcrrrtest_argequivalents  z$TestTakeAlongAxis.test_argequivalentcCstd}tjdtjd}t||ddttt|tdddttt||t ddttt||t ddttj t||dddS)z4 Test it errors when indices has too few dimensions ) r@)r@r+dtyperr@N) ronesintprr ValueErrorarray IndexErrorastypeboolfloat AxisErrorr:rairrr test_invalid1s zTestTakeAlongAxis.test_invalidcCs:td}tjdtjd}t||dd}t|j|jdS)zC Test everything is ok with empty results, even with inserted dims )r(r)r*)r(rr*rArCrNrrDrErrshaper:rrNactualrrr test_emptyBs zTestTakeAlongAxis.test_emptycCs8td}tjdtjd}t||dd}t|jddS)D Test that non-indexing dimensions are broadcast in both directions r(r)rCrCr+r*rArCrr(r+r*NrPrRrrrtest_broadcastJs z TestTakeAlongAxis.test_broadcastN)__name__ __module__ __qualname__r?rOrTrYrrrrr%s  r%c@eZdZddZddZdS)TestPutAlongAxiscCsztgdgdg}tt|jdgD]%}|}ttj||d}t||d|dttj ||d}t ||qdS)N)r@)<(2ri) rrGr7r8r9copyr#r4rr2r)r:a_baserri_maxi_minrrrtest_replace_maxSs z!TestPutAlongAxis.test_replace_maxcCsLtd}tjdtjddd}t||dddtt||dddd S) rUrVr@rArWr)r`rCrN)rrDarangerEreshaperrrrMrrrrYcs zTestPutAlongAxis.test_broadcastN)rZr[r\rhrYrrrrr^Rs r^c@seZdZddZddZddZddZd d Zej fd d Z ej fd dZ ej fddZ ddZ ddZddZddZddZdS)TestApplyAlongAxiscC6tdd}tttd|t|t|jddS)Nr`r@drrCrrDrrlenrQr:rrrr test_simplel  zTestApplyAlongAxis.test_simplecCrl)N)r@ernrrCrorqrrrtest_simple101qrsz!TestApplyAlongAxis.test_simple101cCs:tdd}tttjd|gdgdgdgdS)Nr(r(r(r)rvr_!)$'*)-03)rrirjrrsumrqrrrtest_3dvszTestApplyAlongAxis.test_3dcCsdd}Gdddtj}tddgddgg|}tddgd d gg|}t|d|}tt||t||t|d|}tt||t||dS) NcSs|dS)Nr+r)rowrrrdouble|sz9TestApplyAlongAxis.test_preserve_subclass..doublec@ eZdZdS)z.MyNDArrayNrZr[r\rrrr MyNDArrayrrrCr+r(r))rndarrayrGviewrr isinstancer)r:rrmexpectedresultrrrtest_preserve_subclass{s   z)TestApplyAlongAxis.test_preserve_subclasscCsJGdddtj}dd}td|}tt|d|tgddS)Nc@eZdZdZdS)z9TestApplyAlongAxis.test_subclass..MinimalSubclassrCN)rZr[r\datarrrrMinimalSubclassrcSs|jSN)r)rGrrrminimal_functionsz:TestApplyAlongAxis.test_subclass..minimal_functionrr(r)rCrCrC)rrzerosrrrrG)r:rrrrrr test_subclasss z TestApplyAlongAxis.test_subclasscCsJtd|}ttjd|}tt||t|tgd|dS)Nrrrrr) rrDrrrrrrrG)r:clsrrrrrtest_scalar_arraysz$TestApplyAlongAxis.test_scalar_arraycCsdd}td|}t|d|}tt||t|tgd|t|d|}tt||t|tgd|dS)NcSs t|jdttj|ddS)z/ Sum x, returning a 0d array of the same class rCT)keepdims)rr9rsqueezerxrrr sum_to_0ds z3TestApplyAlongAxis.test_0d_array..sum_to_0drrrrC)r(r(r(r(r(r()rrDrrrrrrG)r:rrrrrrr test_0d_arrays  z TestApplyAlongAxis.test_0d_arraycs2fddtddtd}tjfddtjdDd d }tt |t |t||td}tjfd dtjdDdd }tt |t |t||td d td}tjfddtjdDd d }tt |t |t||dS)Ncs0t|jd|ddd|dddfS)/produces an asymmetric non-square matrix from xrCN)rr9rr)rrrf1to2s $z5TestApplyAlongAxis.test_axis_insertion..f1to2rrcs g|] }dd|fqSrr.0ia2drrr z:TestApplyAlongAxis.test_axis_insertion..rCrrcs g|] }|ddfqSrrrrrrrrZ)rr*r(cs6g|]tjfddtjdDddqS)cs"g|] }|ddfqSrrr)a3drjrrrszETestApplyAlongAxis.test_axis_insertion...rr)rstackr8rQ)r)rr)rrrs r+) rrirjrrr8rQrrtype)r:rrSrr)rrrrrtest_axis_insertionsB         z&TestApplyAlongAxis.test_axis_insertioncCs4Gdddtj}||||||dS)Nc@r)zFTestApplyAlongAxis.test_subclass_preservation..MinimalSubclassNrrrrrrrr)rrrrr)r:rrrrtest_subclass_preservations  z-TestApplyAlongAxis.test_subclass_preservationcCsdd}tdd}t|d|}tt|tjjt|j dt |dddddfj ||dddfj t |dddddfj ||dddfj t |dddddfj ||dddfj dS) NcSs@t|jd|ddd|dddf}tj|ddk|S)rrCNrr*r)rr9rma masked_where)rrrrrrs z8TestApplyAlongAxis.test_axis_insertion_ma..f1to2rrrr(rCr+) rrirjrrrr masked_arrayrr9rmask)r:rrrrrrtest_axis_insertion_mas  004z)TestApplyAlongAxis.test_axis_insertion_mac CsHdd}t|dtddgddgg}t|tddgddggdS)NcSs|d|dfS)NrCrrrrrr sample_1dz7TestApplyAlongAxis.test_tuple_func1d..sample_1drCr+r(r))rrrGr)r:rrrrrtest_tuple_func1ds  z$TestApplyAlongAxis.test_tuple_func1dcCs|dd}td}tttj|d|tttj|d|dd}td}t|d|}t|td tttj|d|dS) NcSs tddS)NF)rrrrr never_calls z1TestApplyAlongAxis.test_empty..never_call)rrrrCcSstt|dkdS)NrrC)rrprrrr empty_to_1sz1TestApplyAlongAxis.test_empty..empty_to_1)r@rr@)remptyrrFrrrD)r:rrrrSrrrrTs  zTestApplyAlongAxis.test_emptycCstddhddhddhgddhd d hd d hgg}td dd|}thdhdhdg}t||t|jD]}tt||t||q;dS)NrC r+r(rxr),r*7rBcSs tj|Sr)setunion)rrrrs z>TestApplyAlongAxis.test_with_iterable_object..r>rCr)rr>r+r*rr>r(rrxr)rrGrrndindexrQr)r:rnrSrrrrrtest_with_iterable_objects z,TestApplyAlongAxis.test_with_iterable_objectN)rZr[r\rrrurrrrrrrrrrrrTrrrrrrkks % rkc@eZdZddZdS)TestApplyOverAxescCsHtdddd}ttj|ddg}t|tdgdgdgggdS) Nr+r(r)rra\|)rrirjrrrrG)r:raoa_arrrrrs"zTestApplyOverAxes.test_simpleN)rZr[r\rrrrrrr rc@4eZdZddZddZddZddZd d Zd S) TestExpandDimscCsRd}t|}tddD]}t||}t|j|dktt|j|kq dS)Nr+r(r)r*r)rC)rrr8r rrQr)r:srrbrrrtest_functionalitys  z!TestExpandDims.test_functionalitycCsntd}tj|ddjdksJtj|ddjdksJtj|ddjdks)Jtj|d djd ks5JdS) NrwrrCr+r)rCrCrCr(r(r()rr)rCr(r(r(rCrC)rr(r*)rCr(r(rCr(rC)rr)rCrCr(rCr(r()rrr rQrqrrrtest_axis_tuple!s zTestExpandDims.test_axis_tuplecCs\d}t|}ttjt|dttjt|dtd}ttjt|dttjt|ddS)Nrr*rw)rr)rr*)rrrrLr )r:rrrrrtest_axis_out_of_range(s  z%TestExpandDims.test_axis_out_of_rangecCstd}ttt|dddS)Nrw)rCrCr)rrrrFr rqrrrtest_repeated_axis2s z!TestExpandDims.test_repeated_axiscCsftdd}tjj||ddkd}tj|dd}tt|tjjt |j dt |j j ddS) Nr@)r+r*r(r)rrCr)r+rCr*) rrirjrrGr rr MaskedArrayrrQr)r:rexpandedrrrtest_subclasses6s  zTestExpandDims.test_subclassesN)rZr[r\rrrrrrrrrrs  rc@sfeZdZddZddZddZddZd d Zej j e d d d dZ ddZ ddZddZdS)TestArraySplitcCtd}ttt|ddS)Nr@r)rrirrFrrqrrrtest_integer_0_splitA z#TestArraySplit.test_integer_0_splitc Cstd}t|d}tdg}t||t|d}tdtddg}t||t|d}tdtddtddg}t||t|d}tdtddtdd td dg}t||t|d}tdtddtddtdd td dg}t||t|d}tdtddtddtdd td d td dg}t||t|d}tdtddtddtddtdd td d td dg}t||t|d }tdtddtddtddtddtdd td d td dg}t||t|d }tdtddtddtddtddtddtdd td d td dg }t||t|d}tdtddtddtddtddtddtddtdd td d td dg }t||t|d }tdtddtddtddtddtddtddtdd td d td dtgg }t||dS) Nr@rCr+r*r(r)r r)rrircompare_resultsrGr:rrdesiredrrrtest_integer_splitEsr                           z!TestArraySplit.test_integer_splitcCsttdtdg}t|ddd}ttdgttdgtdg}t||t|jj|djjut|gddd}tdttdgttdgg}t||t|jj|djjudS)Nr@r(rrrr@rr rrGrirrrrrBrr:rrtgtrrrtest_integer_split_2D_rowss   z)TestArraySplit.test_integer_split_2D_rowsc Csttdtdg}t|ddd}ttdtdgttddtddgttddtddgg}t||dS)Nr@r(rrr)r)rrGrirrrrrrtest_integer_split_2D_colssz)TestArraySplit.test_integer_split_2D_colscCsvttdtdg}t|d}ttdgttdgtdg}t||t|jj|djjudS)z2 This will fail if we change default axis r@r(rrNrrrrrtest_integer_split_2D_defaults   z,TestArraySplit.test_integer_split_2D_defaultzNeeds 64bit platform)reasoncCs\tdgd}t|d}tdgd}|gd}tt|D] }t||j||jqdS)Nr)rr+r))i@r+)r broadcast_torr8rprrQ)r:rrchunkrrrrr,test_integer_split_2D_rows_greater_max_int32s  z;TestArraySplit.test_integer_split_2D_rows_greater_max_int32cCsZtd}gd}t||dd}tddtddtddtddg}t||dS) Nr@)rCr*rrrrrCr*r)rrirrr:rindicesrrrrrtest_index_split_simples  z&TestArraySplit.test_index_split_simplecCsXtd}gd}t||dd}tgtddtddtddg}t||dS)Nr@)rr*rrrrr*rrrirrGrrrrrtest_index_split_low_bounds  z)TestArraySplit.test_index_split_low_boundcCshtd}gd}t||dd}tgtddtddtddtgtgg}t||dS)Nr@)rr*rr@ rrrr*rrrrrrtest_index_split_high_bounds z*TestArraySplit.test_index_split_high_boundN)rZr[r\rrrrrpytestmarkskipifIS_64BITrrrrrrrrr@s=  rc@r]) TestSplitcCs8td}t|d}tdtddg}t||dS)Nr@r+r*)rrirrrrrrtest_equal_splits  zTestSplit.test_equal_splitcCr)Nr@r()rrirrFrrqrrrtest_unequal_splitrzTestSplit.test_unequal_splitN)rZr[r\rrrrrrrs rc@r$) TestColumnStackcCtttddSNrC)r TypeErrorr r:rrrtest_non_iterablerz!TestColumnStack.test_non_iterablecCsLtd}td}tddgddgddgg}t||f}t||dS)N)rCr+r()r+r(r)rCr+r(r)rrGr rr:rrrrSrrrtest_1D_arrayss   zTestColumnStack.test_1D_arrayscCsdtdgdgdgg}tdgdgdgg}tddgddgddgg}t||f}t||dS)NrCr+r(r)r r rrrtest_2D_arrayss zTestColumnStack.test_2D_arrayscCDtttddtdDWddS1swYdS)Ncs|]}tdVqdSr(Nrrir_rrr z1TestColumnStack.test_generator..r+)r FutureWarningr r8rrrrtest_generator "zTestColumnStack.test_generatorN)rZr[r\r r r rrrrrrs  rc@s<eZdZddZddZddZddZd d Zd d Zd S) TestDstackcCrr)rrr rrrrr rzTestDstack.test_non_iterablecCs@td}td}t||g}tddggg}t||dSNrCr+rrGr rr:rrrrrrr test_0D_arrays   zTestDstack.test_0D_arraycCsDtdg}tdg}t||g}tddggg}t||dSrrrrrr test_1D_arrays   zTestDstack.test_1D_arraycCsXtdgdgg}tdgdgg}t||g}tddggddggg}t||dSrrrrrr test_2D_arrays  zTestDstack.test_2D_arraycCsNtddg}tddg}t||g}tddgddggg}t||dSrrrrrrtest_2D_array2 s  zTestDstack.test_2D_array2cCr)Ncsrrrrrrrrrz,TestDstack.test_generator..r+)rrr r8rrrrrrzTestDstack.test_generatorN) rZr[r\r rrrr rrrrrrs rc@0eZdZdZddZddZddZdd Zd S) TestHsplit&Only testing for integer splits. cCtttdddSr)rrFrrrrrr zTestHsplit.test_non_iterablecCs8td}z t|dtdWdStyYdSw)NrCr+r)rrGrrrFrqrrrr!s   zTestHsplit.test_0D_arraycCsBtgd}t|d}tddgtddgg}t||dSNrCr+r(r)r+rCr(r)rrGrrrrrrr)s zTestHsplit.test_1D_arraycCsZtgdgdg}t|d}tddgddggtddgddggg}t||dSr&r(rrrrr/s  ,zTestHsplit.test_2D_arrayNrZr[r\__doc__r rrrrrrrr"s  r"c@r!) TestVsplitr#cCr$r)rrFrrrrrr ;r%zTestVsplit.test_non_iterablecCrr)rrGrrFrrqrrrr>rzTestVsplit.test_0D_arraycCs<tgd}z t|dtdWdStyYdSwNr'r+r)rrGrrrFrqrrrrBs  zTestVsplit.test_1D_arraycCsNtgdgdg}t|d}tgdgtgdgg}t||dSNr'r+)rrGrrrrrrrJs   zTestVsplit.test_2D_arrayNr)rrrrr+7s  r+c@r) TestDsplitcCr$r)rrFrrrrrr Tr%zTestDsplit.test_non_iterablecCrrrrGrrFrrqrrrrWrzTestDsplit.test_0D_arraycCs tgd}ttt|ddSr-r/rqrrrr[szTestDsplit.test_1D_arraycCsDtgdgdg}z t|dtdWdSty!YdSwr,)rrGrrrFrqrrrr_s   zTestDsplit.test_2D_arraycCstgdgdggdgdgg}t|d}tddgddggddgddgggtddgddggddgddgggg}t||dSr&)rrGrrrrrr test_3D_arrayhs  $$zTestDsplit.test_3D_arrayN)rZr[r\r rrrr0rrrrr.Rs  r.c@r) TestSqueezecCsddlm}|ddddd}|ddddd}|dddd}tt|t|dtt|t|dtt|t|dd ggg}t|}t|d t|jdtt|tj dS) Nrr&r`r@rC)r`r@r@)r`r@r`rmg?) r-r'rrrrjrr9rr)r:r'rrcrrrr test_basicts     zTestSqueeze.test_basicNrZr[r\r3rrrrr1srr1c@r)TestKroncCsGdddtj}tddg}||j|j|j}ttt||tjttt|||ttt||tjttt|||dS)Nc@r)z*TestKron.test_return_type..myarraygN)rZr[r\__array_priority__rrrrmyarrayrr7r+) rrrDrQrBrrrr )r:r7rrrrrtest_return_typeszTestKron.test_return_typeN)rZr[r\r8rrrrr5rr5c@r$) TestTilecCstgd}ddgddgg}tt|dgdtt|dgdgdgtt|dgdgtt|dgd gd gtt|d ddgddgddgddggtt|dgd gd gd gd gdS) NrrCr+r(r))rrCr+rrCr+r+r+rCr+)rCr+rCr+)r(r)r(r)r+rC)rrGrr r:rrrrrr3s(  zTestTile.test_basiccCs0td}t|d}|d7}t|tddS)Nr*rCr+)rrir rr=rrr(test_tile_one_repetition_on_array_gh4679s  z1TestTile.test_tile_one_repetition_on_array_gh4679cCsLtggg}tggg}t|dj}t|dj}t|dt|ddS)Nr+rX)r+r)r(r+r)rrGr rQr)r:rrr2rnrrrrTs    zTestTile.test_emptyc Csnddlm}gd}gd}|D]$}|dd|d}|D]}t||j}t||}t||} t|| qqdS)Nr)randint))r+r;r<r:)r+r(r+)r(r+))r()r+r()r(r)r()r(r+r()r)r(r+r)r:r@)size)r-r?rrDrBr r r) r:r?repsrQrrrrlargeklargerrrtest_kroncompares    zTestTile.test_kroncompareN)rZr[r\r3r>rTrErrrrr9s   r9c@r)TestMayShareMemoryc Cstd}td}tt||tt||dddtt||dddtt||dddddftt|ddd| tt|ddd| tt|dddddf| tt|dddddf|dS)N)rcra)r_rarrr+rC)rrDrmay_share_memory)r:rnd2rrrr3s  "$&zTestMayShareMemory.test_basicNr4rrrrrFrrFcCs(tt|D] }t||||qdSr)r8rpr)rrrrrrrsr).numpyrr!sysrnumpy.lib.shape_baserrrrrrrr r r r r rr numpy.testingrrrrrmaxsizerr#r%r^rkrrrrrrr"r+r.r1r5r9rFrrrrrs6@  4'('! (