
    i [                       S r SSKrSSKrSSKrSSKrSSK	r
SSKrSSKrSSKrSSKrSSKJr  SSKJr  SSKJr  SSKJr  SSKJr  SSKJ r!  SS	KJ"r#  SS
KJ$r%  SSKJ&r'  SSKJ(r)  SSKJ*r+  SSK,r-SSK.r/SSK.J0r0J1r1  SSK.J2r3  SSK.J4r4  SSK.J5r6  SSK7J8r9  SSK:J;r<  SSK=J>r?  SSK@JArB  SSK@JCrD  SSKEJFrG  SSKHJIrJ  SSKKJLrM  SSKNJOrP  SSKQJRrS  SSKTJUrV  SSKWJXrY  SSKZJ[r[  S r\S!r]S"r^S#r_S$r`S%ra SS&KbJcrd  \[(       a  SSKeJfrg  S' rhS( ri  SS)S*S+\%\j   S,\%\j   S-\j4S. jjrkS/\jS-\j4S0 jrlSS2 jrmS3\jS-S*4S4 jrnS5 roS6 rpS7 rqS8 rrS9 rsS: rt\3R                  4S; jrvSS< jrwSS= jrx SS> jry SS? jrzS@ r{SA r|SSB jr}SC r~SD rSE rSF rSG r\" 5       SH 5       rSI rSJ rSK rSL r SSM\+SN\j4   SO\%\j   S-\j4SP jjrSSQ.SRSSST\+S\/Rd                  4   S-SS4SU jjrSVSSSW\SX\S-\D4SY jr\0S\?R|                  S14SVSSSZ\+\M\!\M   4   S[\S\\%\   SX\S]\S-\+SS\D4   4S^ jjrS_ r " S` Sa5      rSb\9R                  Sc\9R                  Sd\jSe\jS-S4
Sf jrSg\SO\j4Sh jr SSiSSSj\\j\\j\'\   4   4   SO\jSk\jS-S4
Sl jjrSSm jr  SSM\+\jSS4   Sn\jSo\%\   Sp\%\   4Sq jjr   SSrSSSs\%SS   St\%SS   Su\%SS   S-S4
Sv jjrSw\9R                  S-\'\\\4   4Sx jrSy\9GR.                  Sz\S-\#\9GR0                     4S{ jr\PSy\9GR.                  S|\#\9GR0                     S-S4S} j5       rSw\9R                  Sz\S-\9R                  4S~ jrSw\9R                  Sz\S-\9R                  4S jr   SSSSSSSSS\%\#\j      S\%\#\j      S\%\#\j      S-S4S jjrg!   Srd GN= f)z#
Utilities for the entire package.
    N)OrderedDict)deepcopy)	lru_cache)Callable)Dict)Iterable)List)Optional)Tuple)Type)Union)_SPECIFICATION_VERSION_IOS_16_SPECIFICATION_VERSION_IOS_18)ComputeUnit)_logger)proto)mil)mil_convert)load)Builder)Program)NameSanitizer)WeightRandomizer)AbstractGraphPass)block_context_manager)PassPipelineManager)PASS_REGISTRY)Placeholder   )
_HAS_SCIPY.mlmodel
.mlpackagezmodel.mlmodelz
weight.binweightszcom.apple.CoreML)ModelPackagec                 P    [        U [        5      (       a  U R                  5       $ U $ N)
isinstancebytesdecode)xs    W/var/www/html/ai-image-ml/venv/lib/python3.13/site-packages/coremltools/models/utils.py_to_unicoder,   J   s     !Uxxz    c                     [        U R                  5       5      n[        UR                  R                  R
                   Vs/ s H  o3R                  PM     sn5      nU H  nXT;  d  M
  X	 M     g s  snf r&   )listkeysset_specdescriptioninputname)
input_dictmodelinput_dict_keysinpmodel_input_namesks         r+   _remove_invalid_keysr<   Q   s^     :??,-O1H1H1N1NO1N#XX1NOP%  Ps   A0
proto_spec_proto.Model_pb2weights_dirpackage_pathreturnc                    Uc  [         R                  " [        S9n[        R                  R                  U5      (       a?  [        R                  " U5      (       a  [        SU 35      e[        R                  " U5        [        R                  R                  U5      u  p4U[        :w  a  [        S[         SU S35      e[        U5      n[         R                  " [        S9nUR                  U R!                  5       5        UR#                  5         UR%                  UR&                  [(        [*        S5        UR-                  5         Ub  UR/                  U[0        [*        S5        U$ )a$  

Parameters
----------
proto_spec
    The proto spec of the model.

weights_dir
    Copy weights from this path to the ``mlpackage``.

package_path
    Place the created ``mlpackage`` at this path. Error out if this path is a non-empty directory.

Returns
-------
path to the ``mlpackage``.
)suffixz@The package_path is invalid because it's a non-empty directory: z%For an ML Package, extension must be z (not )zCoreML Model SpecificationzCoreML Model Weights)	_tempfilemkdtemp_MLPACKAGE_EXTENSION_ospathexistslistdirFileExistsError_shutilrmtreesplitext	Exception_ModelPackageNamedTemporaryFile_MLMODEL_EXTENSIONwriteSerializeToStringflushsetRootModelr5   _MODEL_FILE_NAME_MLPACKAGE_AUTHOR_NAMEcloseaddItem_WEIGHTS_DIR_NAME)r=   r?   r@   _extpackage	spec_files          r+   _create_mlpackagera   [   s5   ,  ((0DE
xx|$$;;|$$!RS_R`a  	|$XX|,FA
""34H3IPSuTUV
 	
 L)G ,,4FGIOOJ0023OO)9;Q57 OO ""		
 r-   	save_pathc                 V    [         R                  R                  U S[        [        5      $ )za
Gets the filepath to save a protobuf model if it's within a `mlpackage`.

Parameters
----------
Data)rH   rI   joinrY   rX   )rb   s    r+   _get_model_spec_pathrf      s     88==F,BDTUUr-   Fc                    [         R                  R                  U5      u  pESnU(       d  SR                  U[        5      nO>U[
        :X  a  SnO1U[        :X  a  SnO$[        SR                  [        [
        U5      5      eU(       a   SSKJn  UR                  U 5      n U(       a@  [        c  [        S	5      eU R                  S
5      S:X  a  Uc  [        S5      e[!        XUS9  g[#        US5       n	U	R%                  U R'                  5       5        SSS5        g! [         a0  n[        U5        [        R                  " S[        5         SnANSnAff = f! , (       d  f       g= f)aP  
Save a protobuf model specification to file.

Parameters
----------
spec: Model_pb
    Protobuf representation of the model.

filename: str
    File path where the spec is saved.

auto_set_specification_version: bool
    If ``True``, will always try to set specification version automatically.

weights_dir: str
    Path to the directory containing the weights.bin file. This is required
    when the spec has model type ``mlprogram``. If the ``mlprogram`` does not contain
    any weights, this path can be an empty directory.

Examples
--------
.. sourcecode:: python

    coremltools.utils.save_spec(spec, "HousePricer.mlmodel")
    coremltools.utils.save_spec(spec, "HousePricer.mlpackage")
    coremltools.utils.save_spec(
        spec, "mlprogram_model.mlpackage", weights_dir="/path/to/weights/directory"
    )

See Also
--------
load_spec
Fz{}{}Tz#Extension must be {} or {} (not {})r   _MLModelProxyz=Failed to automatic set specification version for this model.Nz0Unable to load libmodelpackage. Cannot save specr   	mlProgramzspec of type mlProgram cannot be saved without the weights file. Please provide the path to the weights file as well, using the 'weights_dir' argument.)r?   r@   wb)rH   rI   rO   formatrS   rG   rP   libcoremlpythonri   auto_set_specification_versionprint	_warningswarnRuntimeWarningrQ   
WhichOneofra   openrT   rU   )
specfilenamern   r?   r5   r^   
is_packageri   efs
             r+   	save_specrz      sA   D !!(+IDJ==+=>	$	$
	"	"
=DDEWYmorstt%	 8 ??ED  B  ??6"k1k6I !F G G 	$hO(D!QGGD**,- "!%  	!HNNO 	$ "!s$   D 2 E
E%&EE
E&
model_pathc                    [         R                  R                  U 5      (       a:  [        c  [	        S5      e[        U 5      R                  5       R                  5       nOU n[        R                  R                  5       n[        US5       nUR                  UR                  5       5        SSS5        U$ ! , (       d  f       U$ = f)a  
Load a protobuf model specification from file (``mlmodel``) or directory (``mlpackage``).

Parameters
----------
model_path: Path to the model from which the protobuf spec is loaded.

Returns
-------
model_spec: Model_pb
    Protobuf representation of the model.

Examples
--------
.. sourcecode:: python

    spec = coremltools.utils.load_spec("HousePricer.mlmodel")
    spec = coremltools.utils.load_spec("HousePricer.mlpackage")

See Also
--------
save_spec
Nz6Unable to load libmodelpackage. Cannot make save spec.rb)rH   rI   isdirrQ   rP   getRootModel_proto	Model_pb2Modelrt   ParseFromStringread)r{   specfileru   ry   s       r+   	load_specr      s    0 xx~~j!! TUU ,99;@@B!!#D	h	QVVX& 
K 
	Ks    B55
Cc                    / nU R                  S5      S:X  aO  / nU R                  R                   H1  nU(       d  [        U5      s  $ UR	                  [        U5      5        M3     U$ U R                  S5      S;   aO  / nU R                  R                   H1  nU(       d  [        U5      s  $ UR	                  [        U5      5        M3     U$ U R
                  R                  (       a  U R
                  R                  nU$ U R                  R                  (       a  U R                  R                  nU$ U R                  R                  (       a  U R                  R                  nU$ )z
Returns a list of neural network layers if the model contains any.

Parameters
----------
spec: Model_pb
    A model protobuf specification.

Returns
-------
[NN layer]
    list of all layers (including layers from elements of a pipeline).

r   pipeline)pipelineClassifierpipelineRegressor)	rs   r   models_get_nn_layersextendneuralNetworklayersneuralNetworkClassifierneuralNetworkRegressor)ru   r   
model_specs      r+   r   r     s%     Fv*,--..J%j11nZ89	 /* M 
	 $O	O--..J%j11nZ89	 / M 
			"	"##** M 
	%	%	,	,--44 M 
	$	$	+	+,,33Mr-   c                     [         R                  " U 5      nSnU HP  n[         R                  " UR                  5       SS9nX$S   R                  5       -  nX$S   R                  5       -  nMR     U$ )N int8)dtype   r   )_npfloat16
fromstringtobytes)fp32_arrraw_fp16r*   fp16	all_bytess        r+   !_fp32_to_reversed_fp16_byte_arrayr   >  sf    {{8$H
ANN4<<>@		q\!!##	q\!!##  Hr-   c                    [         R                  " U 5      S:  d  [         R                  " U 5      S::  a  [        S5      e[        R
                  S:X  a$  [         R                  " U 5      R                  5       $ [        U 5      $ )Ni  i  zZModel cannot be converted as it has weights that cannot be represented in half precision.
little)	r   amaxaminrP   _sys	byteorderr   r   r   )r   s    r+   _fp32_to_fp16_byte_arrayr   H  sf    
xxU"chhx&8F&B 
 	
 ~~!{{8$,,..0::r-   c                     U (       d   e[        U R                  5      S:X  a  g [        U R                  5      U l        U R                  S S 2	 g Nr   )len
floatValuer   float16Value)wps    r+   _wp_to_fp16wpr   V  s9    I2
2==Q.r}}=BO
ar-   c                 V    SSK Jn  U" U S[        R                  R                  5      nU$ )Nr   )_quantize_spec_weights   )!neural_network.quantization_utilsr   _ctr   &_QUANTIZATION_MODE_LINEAR_QUANTIZATION)fp_specr   qspecs      r+   ,_convert_neural_network_spec_weights_to_fp16r   ^  s"    I"7B

0a0abELr-   c                 J    U R                  5       n[        [        U5      5      $ )a  
Utility function to convert a full-precision (float) MLModel to a
half-precision MLModel (float16).

Parameters
----------
full_precision_model: MLModel
    Model which will be converted to half precision. Currently conversion
    for only neural network models is supported. If a pipeline model is
    passed in, then all embedded neural network models embedded within
    will be converted.

Returns
-------
model: MLModel
    The converted half precision MLModel.

)get_spec
_get_modelr   )full_precision_modelru   s     r+   '_convert_neural_network_weights_to_fp16r   e  s#    &  ((*DB4HIIr-   c                     [        U [        R                  R                  5      (       a  U $ [        R                  R                  XS9$ )z(
Utility to get the model and the data.
compute_units)r'   r   r   MLModel)ru   r   s     r+   r   r   |  s8     $

**++zz!!$!DDr-   c                 "   [        U 5      n U(       a  [        S5        [        S5        SnSnUR                  5        H~  u  pg[        U5      n[	        X5        U R                  U5      [        U5         n	Xr   n
X-
  nU(       a  [        SR                  XU5      5        [        [        U5      U5      nX[U-  -   nM     [        U5      [        R                  " U[        U5      -  5      US.nU(       a  [        SR                  U5      5        U$ )a  
Evaluate a Core ML regression model and compare against predictions
from the original framework (for testing correctness of conversion).

Parameters
----------
model: MLModel or str
    A loaded MLModel or a path to a saved MLModel.

data: Dataframe
    Test data on which to evaluate the models.

target: str
   Name of the column in the dataframe to be compared against the prediction.

verbose: bool
   Set to true for a more verbose output.

See Also
--------
evaluate_classifier

Examples
--------
.. sourcecode:: python

    metrics = coremltools.utils.evaluate_regressor(
        spec, "data_and_predictions.csv", "target"
    )
    print(metrics)
    {"samples": 10, "rmse": 0.0, max_error: 0.0}
r   z!Other Framework		Predicted		Deltar   z{}				{}			{:0.4f})samplesrmse	max_errorresults: {})r   ro   iterrowsdictr<   predictr,   rl   maxabsr   _mathsqrt)r7   datatargetverboser   error_squaredr]   rowr6   	predictedother_frameworkdeltarets                r+   evaluate_regressorr     s    B uEb	56IM--/#Y
Z/MM*-k&.AB	++-44_QVWXE
I.	%7 " t9

=3t945C m""3'(Jr-   c                    [        U 5      n U(       a  [        S5        [        S5        SnUR                  5        Hi  u  pV[        U5      n[	        Xp5        U R                  U5      [        U5         nXb   n	X:w  a  US-  nU(       d  MO  [        SR                  X5      5        Mk     [        U5      US.n
U(       a  [        SR                  U
5      5        U
$ )ab  
Evaluate a Core ML classifier model and compare against predictions
from the original framework (for testing correctness of conversion).
Use this evaluation for models that don't deal with probabilities.

Parameters
----------
filename: list of str or list of MLModel
    File to load the model from, or a loaded
    version of the MLModel.

data: list of str or list of Dataframe
    Test data on which to evaluate the models (dataframe,
    or path to a CSV file).

target: str
   Column to interpret as the target column.

verbose: bool
   Set to true for more verbose output.

See Also
--------
evaluate_regressor, evaluate_classifier_with_probabilities

Examples
--------
.. sourcecode:: python

    metrics = coremltools.utils.evaluate_classifier(
        spec, "data_and_predictions.csv", "target"
    )
    print(metrics)
    {"samples": 10, num_errors: 0}
r   Other Framework		Predictedr   r   z{}				{}num_samples
num_errorsr   )	r   ro   r   r   r<   r   r,   rl   r   )r7   r   r   r   r   r]   r   r6   r   r   r   s              r+   evaluate_classifierr     s    H uEb	,-J--/#Y
Z/MM*-k&.AB	+'!OJ7.''CD " d):
>Cm""3'(Jr-   c           	      x   [        U 5      n U(       a  [        S5        [        S5        Su  pEUR                  5        GHB  u  pg[        U5      R	                  5        VV	s0 s H  u  pX:w  d  M  X_M     n
nn	[        X5        U R                  U
5      [        U5         nXr   n[        UR                  5       5      [        UR                  5       5      :w  aI  U(       a;  [        S[        UR                  5       5      [        UR                  5       5      5        US-  nM  UR	                  5        H6  u  pXU   -
  nU(       a  [        XX   5        [        [        U5      U5      nM8     U(       d  GM7  [        S5        GME     [        U5      UUS.nU(       a  [        SR                  U5      5        U$ s  sn	nf )a  
Evaluate a classifier specification for testing.

Parameters
----------
filename: [str | Model]
    File to load the model from, or a loaded
    version of the MLModel.

data: [str | Dataframe]
    Test data on which to evaluate the models (dataframe,
    or path to a CSV file).

probabilities: str
   Column to interpret as the probabilities column.

verbose: bool
   Verbosity levels of the predictions.
r   r   )r   r   zDifferent classes: r   )r   max_probability_errornum_key_mismatchr   )r   ro   r   r   itemsr<   r   r,   r1   r0   strr   r   r   rl   )r7   r   probabilitiesr   r   r   r]   r   r;   vr6   predicted_valuesother_values	cur_classcur_predicted_class_valuesr   r   s                    r+   &evaluate_classifier_with_probabilitiesr     s   . uEb	,-.2+--/'+Cy'8O'8tqA<Ndad'8
OZ/ ==4[5OP)$$&'3|/@/@/B+CC)(--/0))+,
 !5E5K5K5M1I.i1HHEe9PQ$'E
4I$J! 6N 7"I1 "6 4y!6,C m""3'(JE Ps   F6.F6c                 	   U(       d  U(       d  gSnSnU(       a7  U R                   R                   H  nUR                  U:X  d  M  X'l        SnM     U(       a  U R                   R                   H  nUR                  U:X  d  M  X(l        SnM     U R                   R                  U:X  a  X R                   l        U R                   R
                  U:X  a  X R                   l        U(       d  U(       d  gSn	S H&  n
U R                  U
5      (       d  M  [        X
5      n	M(     U	b  U	R                   H{  nU(       d  M  [        UR                  5       H  u  pX:X  d  M  X+R                  U'   M     U(       d  MH  [        UR                  5       H  u  pX:X  d  M  X+R                  U'   M     M}     U(       a+  U	R                   H  nUR                  U:X  d  M  X.l        M     U R                  S5      (       a  U	R                  U:X  a  X)l        U R                  S5      (       a>  U(       a7  U R                  R                   H  nUR                  U:X  d  M  X'l        SnM     SnU R                  S5      (       a  U R                   nOYU R                  S5      (       a  U R"                  R                   nO,U R                  S	5      (       a  U R$                  R                   nUbc  [        UR&                  5       HJ  u  nn[)        UUUU=(       d    US
:g  U=(       d"    U[+        U R                   R&                  5      :  5        ML     U R                  S5      (       Ga  [-        5       R/                  U5      nUU:w  a  [1        SR3                  U5      5      eU R4                  nUR6                  R9                  5        GH+  nUR:                   H  nUR                  U:X  d  M  UUl        M     UR<                  R9                  5        H  n[        UR>                  5       H  u  nnUU:X  d  M  UUR>                  U'   M     UR@                   H  nUR:                  R9                  5        HG  nURB                   H4  nUR                  S5      (       d  M  UR                  U:X  d  M-  UUl        M6     MI     UR>                   H  nUR                  U:X  d  M  UUl        M     M     M     GM.     gg)a  
Rename a feature in the specification.

Parameters
----------
spec: Model_pb
    The specification containing the feature to rename.

current_name: str
    Current name of the feature. If this feature doesn't exist, the rename
    is a no-op.

new_name: str
    New name of the feature.

rename_inputs: bool
    Search for ``current_name`` only in the input features (that is, ignore output
    features).

rename_outputs: bool
    Search for ``current_name`` only in the output features (that is, ignore input
    features).

Examples
--------
.. sourcecode:: python

    # In-place rename of spec
    model = MLModel("model.mlmodel")
    spec = model.get_spec()
    coremltools.utils.rename_feature(spec, "old_feature", "new_feature_name")
    # re-initialize model
    model = MLModel(spec)
    model.save("model.mlmodel")

    # Rename a spec when the model is an mlprogram, in that case, weights are stored outside of the spec
    model = coremltools.convert(torch_model, convert_to="mlprogram")
    spec = model.get_spec()
    # print info about inputs and outputs
    print(spec.description)
    coremltools.utils.rename_feature(spec, "old_feature", "new_feature_name")
    # re-initialize model
    model = MLModel(spec, weights_dir=model.weights_dir)
    model.save("model.mlpackage")
NFT)r   r   r   r   featureVectorizerr   r   r   r   rj   zInput/output names for ML Program must be of the format [a-zA-Z_][a-zA-Z0-9_]*. That is, it must start with a letter and only contain numerals, underscore or letters. Provided feature name, "{}" does not satisfy these requirements.r5   )"r3   r4   r5   outputpredictedFeatureNamepredictedProbabilitiesNameHasFieldgetattrr   	enumeratepreprocessingfeatureNamelabelProbabilityLayerNamer   	inputListinputColumnr   r   r   r   rename_featurer   _NameSanitizersanitize_name
ValueErrorrl   rj   	functionsvaluesinputsblock_specializationsoutputs
operations	arguments)ru   current_namenew_namerename_inputsrename_outputschanged_inputchanged_outputr4   r   nnnn_typelayerindexr5   preprocess_paramsr   r7   new_name_sanitizedr   functionname_value_typeblockiout_nameopargumentbindings                              r+   r   r   I  s   b MN%%++Ezz\)%
 $ ,
 &&--F{{l*&!% .
 00L@4<166,F:B7 
B
 ==!!'B 
~YYE}#,U[[#9KE+-5E* $: ">'0'>/2:LL/ (?  %'%5%5!$00L@4<1 &6 ==233++|;/7, }}())m++55E  L0$,! $ 6 H}}Z  ==	+	,	,**33	*	+	+))22%hoo6LE5-%1*E53t}}/C/C+D#D 7 }}[!!+-;;HE)) bbhbhiqbrt t nn,,.H#+??"''<7+3O( $3 "77>>@#,U]]#;KAx</+3a( $<  **B$&II$4$4$6'/'9'9G&//77#*<<<#?3;GL (: %7
 ,.::*//<?3;O0 ,6 +	 A	 / "r-   c                 &   [        U [        [        [        45      (       a  U $ [        (       a+  [
        R                  " U 5      (       a  U R                  5       $ [        U [        R                  5      (       a  U $ [        U [        5      (       a	  S U  5       $ [        U [        5      (       a  U  Vs/ s H  n[        U5      PM     sn$ [        U [        5      (       a   [        S U R                  5        5       5      $  [        U 5      5       es  snf )z]
Performs cleaning steps on the data so various type comparisons can
be performed correctly.
c              3   8   #    U  H  n[        U5      v   M     g 7fr&   _sanitize_value).0r   s     r+   	<genexpr>"_sanitize_value.<locals>.<genexpr>  s     .Aq""As   c              3   R   #    U  H  u  p[        U5      [        U5      4v   M     g 7fr&   r  )r  r;   r   s      r+   r  r    s!     S_Q');<s   %')r'   r   intfloatr    _spissparsetodenser   ndarraytupler/   r  r   r   )r*   r   s     r+   r  r    s    
 !c3'((	Qyy{	As{{	#	#	Au		.A..	At		,-.Aq"A..	At		SSSSc!fu	 /s   1Dc                   ^ ^ [        T [        R                  5      (       d  [        T[        R                  5      (       aH   [        [        R                  " T 5      [        R                  " T5      -
  5      S:  R                  5       $ [        T [        5      (       aj  [        T[        5      =(       aS    [        T R                  5       TR                  5       5      =(       a$    [        U U4S jT R                  5        5       5      $ [        T [        5      (       a)  [        T T-
  5      S[        T 5      [        T5      -   -  :  $ [        T [        [        45      (       a  T T:H  $ [        T T:H  5      $ !    g= f)z3
Performs a robust equality test between elements.
gh㈵>Fc              3   H   >#    U  H  n[        TU   TU   5      v   M     g 7fr&   )_element_equal)r  r;   r*   ys     r+   r  !_element_equal.<locals>.<genexpr>
  s#     A1N1Q41..s   ")r'   r   r"  r   asarrayallr   r&  r0   r  r/   r#  bool)r*   r'  s   ``r+   r&  r&    s    !S[[!!Z3;;%?%?	AQ784?DDFF 
At		q$ Bqvvx2BAAA	

 
Au		1q5zDCFSVO444	Ae}	%	%AvAF|	s   AE4 4E8c                   ^ [        U 5      n U(       a!  [        U 5        [        S5        [        S5        Sn[        U5       GH  u  pV[        U[        5      (       d   e[        U5      n[        X%   5      nU(       a*  [        S[        U5      5        [        S[        U5      5        [        U R                  U5      5      m[        U[        5      (       d   e[        T[        5      (       d   e[	        U4S jUR                  5        5       5      n	U(       a  [        S[        U	5      5        [        X5      (       a  M  US-  nGM     [        U5      US	.n
U(       a  [        S
R                  U
5      5        U
$ )a  
Evaluate a transformer specification for testing.

Parameters
----------
model: list of str or list of MLModel
    File to load the Model from, or a loaded
    version of the MLModel.

input_data: list of dict
    Test data on which to evaluate the models.

reference_output: list of dict
    Expected results for the model.

verbose: bool
    Verbosity levels of the predictions.

Examples
--------
.. sourcecode:: python

    input_data = [{"input_1": 1, "input_2": 2}, {"input_1": 3, "input_2": 3}]
    expected_output = [{"input_1": 2.5, "input_2": 2.0}, {"input_1": 1.3, "input_2": 2.3}]
    metrics = coremltools.utils.evaluate_transformer(
        scaler_spec, input_data, expected_output
    )

See Also
--------
evaluate_regressor, evaluate_classifier
r   r   r   zInput:
	zCorrect output:
	c              3   0   >#    U  H  oTU   4v   M     g 7fr&    )r  r;   r   s     r+   r  'evaluate_transformer.<locals>.<genexpr>I  s      LOqYq\!2Os   zPredicted:
	r   r   r   )r   ro   r   r'   r   r  r   r   r0   r&  r   rl   )r7   
input_datareference_outputr   r   r
  r   sanitized_rowref_datapredicted_trimmedr   r   s              @r+   evaluate_transformerr5    s/   B uEeb	,-J
+
#t$$$$',"#3#:;,C)'X7#EMM-$@A	(D)))))T****  LHMMO LL"C(9$:;/::!OJ' ,* j/
DCm""3'(Jr-   c                 Z    [        U 5      nU H  nUR                  S5      S:X  d  M    g   g)a  

Returns true if the given protobuf specification has a custom layer, and false otherwise.

Parameters
----------
spec: mlmodel spec

Returns
-------

``True`` if the protobuf specification contains a neural network with a custom layer, ``False`` otherwise.

r	  customTF)r   rs   )ru   r   r	  s      r+   _has_custom_layerr8  X  s3      D!FG$0  r-   c                     [        U 5      n[        5       nU H?  nUR                  S5      S:X  d  M  UR                  UR                  R
                  5        MA     U$ )z

Returns a list of ``className`` fields which appear in the given protobuf spec.

Parameters
----------
spec: mlmodel spec

Returns
-------
set(str)
    A set of unique ``className`` fields of custom layers that appear in the model.

r	  r7  )r   r1   rs   addr7  	classNameru   r   
layers_outr	  s       r+   _get_custom_layer_namesr>  p  sP     D!FJG$0NN5<<112  r-   c                     [        U 5      n/ nU H+  nUR                  S5      S:X  d  M  UR                  U5        M-     U$ )z

Returns a list of all neural network custom layers in the spec.

Parameters
----------
spec: mlmodel spec

Returns
-------
[NN layer]
    A list of custom layer implementations.
r	  r7  )r   rs   appendr<  s       r+   _get_custom_layersrA    sF     D!FJG$0e$  r-   c                     [        U 5      nU H/  nUR                  R                  U:X  d  M  X$R                  l        M1     g)a  

Substitutes ``newname`` for ``oldname`` in the ``className`` field of custom layers. If there are no custom layers, or no
layers with ``className`` = ``oldname``, then the spec is unchanged.

Parameters
----------
spec: mlmodel spec

oldname: str
    The custom layer ``className`` to be replaced.

newname: str
    The new ``className`` value to replace ``oldname``.

Returns
-------
An mlmodel spec.

N)rA  r7  r;  )ru   oldnamenewnamer   r	  s        r+   _replace_custom_layer_namerE    s4    *  %F<<!!W,%,LL" r-   c                  (    [         R                  S:H  $ )z;Returns True if current platform is MacOS, False otherwise.darwin)r   platformr.  r-   r+   	_is_macosrI    s    ==H$$r-   c                  P   [        5       (       a   [        R                  " SS/[        R                  S9R                  R                  S5      R                  S5      n [        U R                  S5       Vs/ s H  n[        U5      PM     sn5      $ gs  snf !   [        S5      e= f)	z
Returns macOS version as a tuple of integers, making it easy to do proper
version comparisons. On non-Macs, it returns an empty tuple.
sw_versz-productVersion)stdoutzutf-8
.z%Unable to determine the macOS versionr.  )rI  _subprocessrunPIPErL  r)   stripr#  splitr  rP   )ver_strr   s     r+   _macos_versionrU    s     {{	E!ooy2C&D[M]M]^eellmtu{{  }A  BG'--*<=*<Q#a&*<=>>  >	ECDDs   A%B 6BB B B%c            
          [         R                  R                  S5      S   n [        [	        [
        [        U R                  S5      5      5      5      n [        U 5      $ )z.
Return python version as a tuple of integers
 r   rN  )r   versionrS  r/   mapr  r#  )rX  s    r+   _python_versionrZ    sF     ll  %a(G3sDs!3456G>r-   c                     U R                   R                   H  nUR                  U:X  d  M  Us  $    U R                   R                   H  nUR                  U:X  d  M  Us  $    [	        SR                  U5      5      e)Nz#Feature with name {} does not exist)r3   r4   r5   r   rP   rl   )ru   feature_nameinput_featureoutput_features       r+   _get_featurer_    sq    ))//-   0 **11,.!! 2 9@@N
OOr-   c                 n    U R                   R                   Vs/ s H  oR                  PM     nnU$ s  snf )z
Returns a list of the names of the inputs to this model.
:param spec: The model protobuf specification
:return: list of str A list of input feature names
)r3   r4   r5   )ru   featureretvals      r+   _get_input_namesrc    s3     +/*:*:*@*@A*@wll*@FAM Bs   2c                 d   S nU R                   R                   H  nU" U5        M     U R                   R                   H  nU" U5        M     U R                   R                   H  nU" U5        M     U R	                  S5      S:X  a)  U R
                  R                   H  n[        U5        M     gg)a  
Convert all double multiarrays feature descriptions (input, output, training input)
to float multiarrays.

Parameters
----------
spec: Model_pb
    The specification containing the multiarrays types to convert.

Examples
--------
.. sourcecode:: python

    # In-place convert multiarray type of spec
    spec = mlmodel.get_spec()
    coremltools.utils.convert_double_to_float_multiarray_type(spec)
    model = coremltools.models.MLModel(spec)
c                 N   U R                   R                  S5      (       a  U R                   R                  R                  [        R
                  R                  R                  :X  a>  [        R
                  R                  R                  U R                   R                  l        g g g )NmultiArrayType)	typer   rf  dataTyper   r   ArrayFeatureTypeDOUBLEFLOAT32)ra  s    r+   _convert_to_floatBconvert_double_to_float_multiarray_type.<locals>._convert_to_float  sq    <<  !122||**33v7G7G7X7X7_7__7=7G7G7X7X7`7`++4 ` 3r-   r   r   N)r3   r4   r   trainingInputrs   r   r   'convert_double_to_float_multiarray_type)ru   rl  ra  r   s       r+   ro  ro    s    (a
 ##))'" * ##**'" + ##11'" 2 v*,--..J3J? / -r-   r7   _proto.Model_pb2.Modeldestination_pathc                 2   S[         S[        4S jn[        5       S:  a  [        S5      e SSKJn  S	n[        U [         5      (       a  U" U 5      (       a  U nO[        S
U  S35      eUch  [        U [        R                  R                  5      (       a  [        S5      e[        U [        R                  R                  5      (       d  [        S5      eUb0  UR                  S5      R!                  S5      (       d  [        S5      eUc@  ["        R$                  " 5        nUS-   n['        X5        UR)                  U5      nS	S	S	5        OUR)                  U5      nUc  W$ [*        R,                  " WU5        U$ !   [        S5      e= f! , (       d  f       N<= f)a  
Compiles a Core ML model.

Parameters
----------
model: Model_pb2 | str
    Either a Core ML model specification (protobuf object) or a path to a saved model (.mlmodel or .mlpackage) file.

destination_path: str
    Path where the compiled model will be saved.

Returns
-------

str : Path to compiled model directory
    If the ``destination_path`` is specified, that is the value that will be returned.

Examples
--------
.. sourcecode:: python

    from coremltools.models import CompiledMLModel
    from coremltools.models.utils import compile_model
    from coremltools.proto import Model_pb2

    spec = Model_pb2.Model()
    spec.specificationVersion = 1

    input_ = spec.description.input.add()
    input_.name = "x"
    input_.type.doubleType.MergeFromString(b"")

    output_ = spec.description.output.add()
    output_.name = "y"
    output_.type.doubleType.MergeFromString(b"")
    spec.description.predictedFeatureName = "y"

    lr = spec.glmRegressor
    lr.offset.append(0.1)
    weights = lr.weights.add()
    weights.value.append(2.0)

    compiled_model_path = compile_model(spec)
    model = CompiledMLModel(compiled_model_path)
    y = model.predict({"x": 2})

See Also
--------
coremltools.models.CompiledMLModel
source_pathrA   c                 v    U R                  S5      nUR                  S5      =(       d    UR                  S5      $ )N/r"   r!   )rstripendswith)rs  	extensions     r+   is_valid_model_file*compile_model.<locals>.is_valid_model_fileJ  s3    &&s+	!!,/Q93E3Ej3QQr-   )
      zDCompiling a Core ML models is only support on macOS 10.13 or higher.r   rh   z%Unable to compile any Core ML models.NzInvalid model file path: 'z7'.Expected a file with .mlpackage or .model extension. z_This model has already been compiled. Call "get_compiled_model_path" to get the compiled model.zQUnrecognized input for "model" parameter. It should be a spec or model file path.ru  z	.mlmodelczB"destination_path" parameter must have ".mlmodelc" file extension.z/spec.mlmodel)r   r+  rU  rP   rm   ri   r'   r   r   r   r   	TypeErrorr   r   r   rv  rw  rE   TemporaryDirectoryrz   compileModelrM   move)r7   rq  ry  ri   source_model_pathsave_dirspec_file_pathoriginal_compiled_model_paths           r+   compile_modelr    s   lR R R
 ("^__A3 %u%% %,UG 4H H 
  eSZZ//00. 
 %!1!1!7!788c 
 #,<,C,CC,H,Q,QR],^,^`aa  ))+x%7Ne,+8+E+En+U( ,+
 (5'A'ABS'T$ ++LL-/?@YA?@@> ,+s   E8 "F8F
Fr   r   _ct.models.MLModelr   c                 	  ^^ U4S jm[        T5      S:  d   eU b   [        U [        5      (       d  [        S5      eU c=  [	        [        U4S jTSS 5      5      nU(       d  [        S5      eTS   R                  n U [        R                  :X  a(  [        5       (       a  [        5       S:  a  [        S	5      e[        [        S
 T5      5      n[        R                  R                  R                  5       n[!        [        S U5      5      Ul        [%        5       nU H  nUR&                  R(                   Hc  nUR*                  U;  d  M  UR&                  R(                  R-                  5       R/                  U5        UR-                  UR*                  5        Me     UR1                  UR&                  R2                   Vs/ s H  oR*                  PM     sn5        M     [%        5       n	USSS2    H  nUR&                  R2                   HH  n
U
R*                  U	;  d  M  UR&                  R2                  R-                  5       R/                  U
5        MJ     U	R1                  UR&                  R(                   Vs/ s H  oR*                  PM     sn5        M     0 n[5        [        U5      S-
  5       H  nX8S-      R&                  R(                   H  nUR6                  XR*                  '   M     X8   R&                  R2                   H  nUR6                  R9                  S5      S:X  a/  [        UR6                  R:                  R<                  5      S:w  a  MQ  UR*                  U;   d  Mc  UR6                  R?                  XR*                     5        M     M     [A        U5       HX  u  pUR9                  S5      S:X  a  SU S3nT" URB                  U5        URD                  RF                  RI                  U5        MZ     [K        U5      nUS-   [L        -   S-   [N        -   n[P        RR                  " U5        [A        T5       Ho  u  nnURT                  c  M  URT                  S-   [V        -   n[P        RX                  R[                  U5      (       d  MQ  [\        R^                  " UUSU S3-   5        Mq     [        RF                  Ra                  X@US9$ s  snf s  snf )a  
Makes a pipeline with the given models.

Parameters
----------
*models :
    Two or more instances of ``ct.models.MLModel``.

compute_units :
    The set of processing units that all models in the pipeline can use to make predictions.
    Can be ``None`` or ``coremltools.ComputeUnit``.

    * If ``None``, the ``compute_unit`` will be inferred from the ``compute_unit`` values of the models.
      If all models do not have the same ``compute_unit`` values, this parameter must be specified.

    * ``coremltools.ComputeUnit`` is an enum with four possible values:
        - ``coremltools.ComputeUnit.ALL``: Use all compute units available, including the
          neural engine.
        - ``coremltools.ComputeUnit.CPU_ONLY``: Limit the model to only use the CPU.
        - ``coremltools.ComputeUnit.CPU_AND_GPU``: Use both the CPU and GPU,
          but not the neural engine.
        - ``coremltools.ComputeUnit.CPU_AND_NE``: Use both the CPU and neural engine, but
          not the GPU. Available only for macOS >= 13.0.

Returns
-------
ct.models.MLModel

Examples
--------
.. sourcecode:: python

    my_model1 = ct.models.MLModel("/tmp/m1.mlpackage")
    my_model2 = ct.models.MLModel("/tmp/m2.mlmodel")

    my_pipeline_model = ct.utils.make_pipeline(my_model1, my_model2)

    y = my_pipeline_model.predict({"x": 12})

    my_pipeline_model.save("/tmp/my_pipeline.mlpackage")
    new_my_pipeline = ct.model.MLModel("/tmp/my_pipeline.mlpackage")

c                    > [        U 5      [        R                  R                  :X  aC  U R	                  S5      S:X  a-  U R
                  R                  S:X  d   eXR
                  l        g g [        U S5      (       a$  U R                  5        H  nT" US   U5        M     g [        U S5      (       a   U R                  5        H  nT" X15        M     g [        U [        5      (       a(  [        U [        5      (       d  U  H  nT" XA5        M     g g g )NvalueblobFileValuez@model_path/weights/weight.bin
ListFieldsr   r   )rg  r   MIL_pb2Valuers   r  fileNamehasattrr  r   r'   	_Iterabler   )proto_messagenew_pathry   r   rx   updateBlobFileNames        r+   r  )make_pipeline.<locals>.updateBlobFileName  s    &.."6"66''0OC$22;;?____7?++4 D ]L11"--/"1Q42 0]H--"))+"1/ ,y11*]TW:X:X""1/ # ;Y1r-   r   NzI"compute_units" parameter must be None or of type coremltools.ComputeUnitc                 :   > TS   R                   U R                   L $ r   )compute_unit)mr   s    r+   <lambda>make_pipeline.<locals>.<lambda>  s    fQi,,>r-   z[Models have different compute_unit values. The "compute_units" parameter must be specified.r   )r|  r   zEcoremltools.ComputeUnit.CPU_AND_NE is only available on macOS >= 13.0c                 "    U R                  5       $ r&   )r   )r  s    r+   r  r    s
    QZZ\r-   c                     U R                   $ r&   specificationVersion)ru   s    r+   r  r    s
    22r-   r   rf  rj   z@model_path/weights/z-weight.binz/Data/ru  )r   r?   )1r   r'   _ComputeUnitr}  r*  rY  r   r  
CPU_AND_NErI  rU  r/   r   r   r   r   r   r  r1   r3   r4   r5   r:  	MergeFromupdater   rangerg  rs   rf  shapeCopyFromr   rj   r   r   r@  ra   rY   r\   rH   mkdirr?   _WEIGHTS_FILE_NAMErI   rJ   rM   copyfiler   )r   r   all_compute_units_the_sameinput_specspipeline_specavailable_as_inputcur_spec	cur_inputr  used_as_input
cur_outputvar_name_to_typejnew_file_pathmlpackage_pathdst	cur_modelweight_file_pathr  s    `                @r+   make_pipeliner    s   `0( v;?? M<)P)Pcdd%(>12J*
 &" *m  q	..000kk(S
 	
 s16:;KII''--/M),2K@*M& !--33I~~%77))//335??	J"&&y~~6 4 	!!83G3G3N3N"O3Na663N"OP   EM"%"..55Jm3))00446@@L 6 	h.B.B.H.HI.Hff.HIJ	 & 3{#a'(U#//55A'(vvVV$ 6 ++22Avv  (,<<QVVEZEZE`E`AaefAfvv)) 0 89 3	 ) !-v&+521#[AMx11=A%%,,X6	 . '}5N
8
#&<
<s
BEV
VCIIcN "&)9  ,(44s:=OOxx/00  !131QC{9K3KL	 * ::mVYZZW #P Js   S
;S
mlmodelspecification_versionpymil_load_funcc                     U R                  5       nUR                  S5      nUS;   a  Sn[        UR                  U5      5      eUS:X  a  O[        SR                  U5      5      eU" UUU R                  S9nU$ )zD
A utility that converts an ``mlprogram`` model into PyMIL program.
r   )r   r   r   r   PipelineClassifierPipelineRegressorzcoremltools.optimize.coreml are meant to be used only with mlprogram typed coreml models. This model has type {}. Please use coremltools.models.neural_network.quantization_utils.quantize_weightsinstead to compress the weights of the model.rj   z3weight compression not applicable for model type {}r   r  file_weights_dir)r   rs   r}  rl   r?   )r  r  r  r   
model_typemsgprogs          r+   !_convert_model_spec_to_pymil_progr    s     !!#J&&v.J  < 	
 

:.//	{	"MTTU_`aa3 ,,D
 Kr-   
graph_passspec_versionskip_model_loadreturn_pymil_progc           
         Uc  U R                   S L nU R                  5       n[        UR                  U5      n[	        XU5      n[        U[        5      (       a  U/nU H@  n	[        U	[        5      (       d  [        S[        U5       35      eU	R                  U5        MB     U(       a  U$ [        USSUU R                  UR                  US9n
U
$ )Nz:graph pass must be an AbstractGraphPass instance, but got 	mlprogrammilinternal
convert_toconvert_fromr  r   model_descriptionr  )	__proxy__r   r   r  r  r'   _AbstractGraphPassAssertionErrorrg  apply_mil_convertr  r3   )r  r  r  r  r  r  r   r  r  cur_graph_passcompressed_mlmodels              r+   _apply_graph_passr  ?  s     !++t3 !!#J
 ? ?N,W_]D *011 \
$.*<== OPTU_P`Oab  	T" %  &"3**$00' r-   c                     Sn [         R                  " U 5      (       a6  [        U 5      R                  [        [        5      nUb  UR                  5       nU$ !    U$ = f)z
Try to find the weights in mlpackage and return the path to the weights directory if found.
Return None if not found.
:param mlpackage_path: str, path to the mlpackage directory
:return: path to the weights directory inside the mlpackage directory
N)rQ   isValidfindItemByNameAuthorr\   rY   rI   )r  r?   	item_infos      r+   _try_get_weights_dir_pathr  l  sf     K  00%n5JJ!#9I $'nn. s   AA Ac                       \ rS rSrSrSS\\   4S jjrS\\\	\\4   4   4S jr
S\SS4S jr\S\\S4   4S	 j5       r\R                  S
\SS4S j5       rS\S\S\SS4S jrS\SS4S jrS\SS4S jrSrg)MultiFunctionDescriptori  a*  
This data class defines how to construct a multifunction model from different model sources.
Use the ``add_function`` method to specify the path to the source ``mlpackage``,
along with the source and target function names.

After setting the ``default_function_name`` to the ``MultiFunctionDescriptor`` instance,
you can export a multifunction model using the ``save_multifunction`` method.

Examples
--------
.. sourcecode:: python

    from coremltools.utils import MultiFunctionDescriptor, save_multifunction

    # Initialize a MultiFunctionDescriptor instance with functions in an existing mlpackage.
    # desc will contain all functions in "my_model.mlpackage"
    desc = MultiFunctionDescriptor("my_model.mlpackage")

    # Construct a MultiFunctionDescriptor instance from scratch.
    # The below code inserts the "main" function from "my_model.mlpackage" as "main_1",
    # and inserts the "main" function from "my_model_2.mlpackage" as "main_2".
    desc = MultiFunctionDescriptor()
    desc.add_function(
        model_path="my_model.mlpackage",
        source_function_name="main",
        target_function_name="main_1",
    )
    desc.add_function(
        model_path="my_model_2.mlpackage",
        source_function_name="main",
        target_function_name="main_2",
    )

    # Each MultiFunctionDescriptor instance must have a default function name
    # so it can be saved as a multifunction mlpackage on disk.
    desc.default_function_name = "main_1"
    save_multifunction(desc, "my_multifunction_model.mlpackage")

See Also
--------
save_multifunction

Nr{   c                 f    SU l         0 U l        0 U l        0 U l        Ub  U R	                  U5        gg)z
If ``model_path`` is passed to the constructor, it must be a :obj:`str` pointing to an
existing ``mlpackage`` on disk. The :py:class:`MultiFunctionDescriptor` instance will be initiated
with the functions in ``model_path``.
N)_default_function_name_name_to_source_function_modelpath_to_functions_modelpath_to_spec	add_model)selfr{   s     r+   __init__ MultiFunctionDescriptor.__init__  s;     '+#(*%')$"$!NN:& "r-   rA   c                 B    [         R                  " U R                  5      $ )z+
Returns ``self._name_to_source_function``
)_copycopyr  r  s    r+   
_functions"MultiFunctionDescriptor._functions  s     zz$7788r-   c                 |   XR                   ;   a  g [        U5      nUR                  n[        UR                  5      S:X  a  S/U R                   U'   O3UR                   Vs/ s H  oUR                  PM     snU R                   U'   X R                  U'   g! [         a  n[        SU SU S35      eSnAff = fs  snf )zT
Given an ``mlpackage`` path ``model_path``, this function caches related metadata.
Nzinvalid model_path z with error z while loading.r   main)	r  r   rP   r   r3   r   r   r5   r  )r  r{   ru   errdescfuncs         r+   _add_modelpath_to_cache/MultiFunctionDescriptor._add_modelpath_to_cache  s     555	aZ(D  t~~!#8>xD((4NRnn7]nd		n7]D((4.2
+  	a2:,l3%_``	a 8^s   B !B9
B6B11B6c                     U R                   $ r&   )r  r  s    r+   default_function_name-MultiFunctionDescriptor.default_function_name  s    ***r-   valc                 X    [        U[        5      (       d  [        SU S35      eXl        g )Nz/default_function_name must be type of str. Got rN  )r'   r   r   r  )r  r  s     r+   r  r    s+    #s##NseSTUVV&)#r-   src_function_nametarget_function_namec                     U R                  U5        X R                  U   ;  a  [        SU SU S35      eX0R                  ;   a  [        SU S35      eX4U R                  U'   g)z
Insert a ``src_function_name`` function from ``model_path`` as the
``target_function_name`` function in the multifunction descriptor.
zsrc_function_name z not found in rN  z	function z already exist.N)r  r  r   r  )r  r{   r  r  s       r+   add_function$MultiFunctionDescriptor.add_function  sy     	$$Z0$@$@$LL12C1DNS]R^^_`aa#@#@@y)=(>oNOO?I>]%%&:;r-   c                 v    U R                  U5        U R                  U    H  nU R                  XU5        M     g)z
Insert all functions from the model in ``model_path`` into the multifunction descriptor.
The function names will remain the same as in the original model.
N)r  r  r  )r  r{   	func_names      r+   r  !MultiFunctionDescriptor.add_model  s8    
 	$$Z055jAIjY? Br-   function_namec                 Z    XR                   ;  a  [        SU S35      eU R                   U	 g)zH
Remove a function ``function_name`` from the multifunction descriptor.
zfunction_name  not found.N)r  r   )r  r  s     r+   remove_function'MultiFunctionDescriptor.remove_function  s2      = ==~m_KHII))-8r-   )r  r  r  r  r&   )__name__
__module____qualname____firstlineno____doc__	_Optionalr   r  _Dict_Tupler  r  property_Unionr  setterr  r  r  __static_attributes__r.  r-   r+   r  r    s    *X'9S> '9E#vc3h'7"78 93# 3$ 3, +vc4i'8 + + !!* * * "*
^^25^MP^	^"@C @D @9S 9T 9r-   r  multifunction_progunifunction_progsrc_func_nametarget_func_namec                 @    U R                  X1R                  U   5        g r&   )r  r   )r  r  r  r  s       r+   1_multifunction_program_append_unifunction_programr    s     ##$46P6PQ^6_`r-   r  c           
      8   SSS[         S[        R                  R                  4S jn0 nU R                  R                  5        HB  u  pEUS   nXc;   a  M  U R                  U   n[        U5      nUc  [        SU S	35      eXx4X6'   MD     [        [        S
 UR                  5       5      5      n	[        U	[        5      n	0 n
UR                  5        H$  u  nu  px[        R                  " UU	U5      nXU'   M&     [        R                   " 5       n0 nU R                  R                  5        Hd  u  pUS   nUS   nX   n["        R$                  R'                  XX5        X6   S   nU" X5      nUR(                  S:X  d   S5       eUUl        UX'   Mf     [*        S   nUR-                  U5        U R.                  nUc  [        S5      eUUR0                  ;  a2  [        SU S[3        UR0                  R5                  5       5       35      eUUl        / nUR0                   H  nUR7                  UU   5        M     [        R                  R9                  UUS9nSUl        SUl        [?        USSU	["        R@                  RB                  USS9nURE                  U5        g)a  
Save a :py:class:`MultiFunctionDescriptor` instance into a multifunction ``mlpackage``.
This function also performs constant deduplication across functions to allow for weight sharing.

Parameters
----------
desc: MultiFunctionDescriptor
    Multifunction descriptor to save on the disk.

destination_path: str
    The path where the new ``mlpackage`` will be saved.

Examples
--------
.. sourcecode:: python

    from coremltools.utils import MultiFunctionDescriptor, save_multifunction

    desc = MultiFunctionDescriptor("my_model_1.mlpackage")
    desc.add_function("my_model_2.mlpackage", "main", "main_2")
    desc.default_function_name = "main_2"

    save_multifunction(desc, "multifunction_model.mlpackage")

See Also
--------
MultiFunctionDescriptor

ru   r>   r   rA   c                    U R                   n[        UR                  5      S:X  ac  US:X  d
   SU 35       e[        R                  R                  UR                  UR                  UR                  UR                  UR                  S9$ UR                   HM  nUR                  U:w  a  M  [        R                  R                  5       nUR                  U5        SUl        Us  $    g)z@
Utils to construct a FunctionDescription from the source spec.
r   r  zinvalid function name )r4   r   stater   r   r   N)r3   r   r   r   r   FunctionDescriptionr4   r   r  r   r   r5   r  )ru   r   
model_desc	func_descress        r+   get_function_spec-save_multifunction.<locals>.get_function_spec/  s     %%
z##$)&L*@(LL&##77 &&!(( &&%/%D%D+5+P+P 8   $--I~~*""668CLL#CHJ .r-   r   Nzweight_dir for model_path r  c                      U S   R                   $ r   r  )r  s    r+   r  $save_multifunction.<locals>.<lambda>X  s    A33r-   r   r   z&function_spec should not have name setzcommon::const_deduplicationzmdefault_function_name must be set for the MultiFunctionDescriptor instance before calling save_multifunction.zdefault_function_name z9 not found in the program. Available functions names are )r   defaultFunctionNameTr  r  r  )#r   r   r   r  r  r   r  r  r   r   rY  r   r   _milproto_to_pymilr   _milr   r   utilsr  r5   _PASS_REGISTRY#_deduplicate_const_across_functionsr  r   r/   r0   r@  ModelDescriptionskip_all_passesexport_as_multifunctionr  r   CPU_ONLYsave)r  rq  r   modelpath_to_spec_and_weightdirr;   r   r{   ru   
weight_dirr  modelpath_to_pymilr  r  function_to_descr  r  function_specr  r  r   r  r  r  s                          r+   save_multifunctionr4    s   B -0				-	-6 ')#--335qT
8&&z2.z:
9*[QRR7;6H'3 6 35T5[5[5]^L |%BCL *I*O*O*Q&
&T!&&

 *.:& +R #<<BBDqT
!!-		CCm	

 /:1=)$>!!R'Q)QQ'--:*  E$   =>J223EF !66${
 	
 $6$@$@@$%:$;;tuy  {M  {W  {W  {\  {\  {^  v_  u`  a
 	
 0E, I",,)$/0 -((991 :  *.&15."*oo..+G LL!"r-   dynamic_shape_mlmodel$function_name_to_materialization_mapsource_function_namec           	         [        U [        R                  R                  5      (       d  [	        S[        U 5       35      eU R                  R                  R                   H-  nUR
                  R                  S5      S:w  d  M$  [        S5      e   U R                  R                  R                   H-  nUR
                  R                  S5      S:w  d  M$  [        S5      e   U R                  b  U R                  nO@[        R                  " U R                  U R                  R                  U R                   5      n[        R"                  R$                  nUR'                  SS5        UR)                  SUUS	.5        [*        R,                  " Xg5        [/        UR0                  5      S
:X  a"  [/        U5      S
:X  a  X1;   a  US:X  a  SUl        OSUl        U R                  R                  nUR2                  (       a$  [5        [        R6                  R8                  U5      nSUl        [=        USSU[        R>                  R@                  SS9n	U	RC                  U5        g)aY  
Given a dynamic-shape mlmodel, materialize symbols to create fixed-shape functions,
then save as an .mlpackage to destination path.
To save memory, the pymil program of input dynamic-shape mlmodel is re-used.
Constant deduplication across functions is performed to allow weight sharing.

Parameters
----------
dynamic_shape_mlmodel : ct.models.MLModel
    A dynamic-shape mlmodel to be materialized

function_name_to_materialization_map: Dict[str, Dict[str, Tuple[int]]]
    A dictionary specifying the name of new functions to be created,
    and for each new function what is the new fixed shapes for inputs.
    If a new function has the same name as an old function,
    then the old function will be overridden

destination_path : str
    The saved .mlpackage model path

source_function_name: str
    The name of the source symbolic-shape function to be materialized, default = main

Examples
--------
.. sourcecode:: python

    from coremltools.utils import materialize_dynamic_shape_mlmodel

    # A dynamic-shape mlmodel you have converted
    dynamic_shape_mlmodel: ct.models.MLModel

    # As an example, let us assume the inputs are
    # 1. ``input_ids (1, query_length)``
    # 2. ``mask (query_length, context_length)``
    function_name_to_materialization_map = {
        "function_name_to_materialization_map": {
            "materialization_2_3": {"input_ids": (1, 2), "mask": (2, 3)},
            "materialization_4_5": {"input_ids": (1, 4), "mask": (4, 5)},
        }
    }

    materialize_dynamic_shape_mlmodel(
        dynamic_shape_mlmodel,
        function_name_to_materialization_map,
        "materialized_model.mlpackage",
    )

To make prediction from the materialized mlmodel, load the desired materialized function

.. sourcecode:: python

    materialization_2_3 = ct.models.MLModel(
        "materialized_model.mlpackage", function_name="materialization_2_3"
    )
    materialization_4_5 = ct.models.MLModel(
        "materialized_model.mlpackage", function_name="materialization_4_5"
    )

See Also
--------
coremltools.converters.mil.mil.passes.defs.experiment.materialize_symbolic_shape_program

zADynamic shape mlmodel must be type of ct.models.MLModel, but got r   rf  z Only tensor input is handled yetz!Only tensor output is handled yetNr   z*common::materialize_symbolic_shape_program)r6  r7  r   r  FTr  r  )r  r  r  r   r  )"r'   r   r   r   r   rg  r2   r3   r4   rs   NotImplementedErrorr   _mil_programr%  r   r  r?   PassPipelineDEFAULTinsert_passset_options_PassPipelineManagerapply_pipeliner   r   r,  r   r   iOS18r+  r  r   r-  r.  )
r5  r6  rq  r7  r4   r   dynamic_shape_progpass_pipeliner  materialized_mlmodels
             r+   !materialize_dynamic_shape_mlmodelrE    s   L +SZZ-?-?@@1235
 	
 ',,88>>::  (,<<%&HII ? (--99@@;;!!&)-==%&IJJ A ))52??/44!''!''<<!--
 $$,,Ma!MN44X$8	
 ''(:J 	(()Q.45: H F*5:2592
 277LL11 #CJJ$4$46K L)-&'"3oo.. ./r-   c                 ,    [        U [        5       SS9nU$ )a~  
Utility function to randomize weights

Parameters
----------
mlmodel: MLModel
    Model which will be randomized.

Returns
-------
model: MLModel
    The MLModel with randomized weights.

Examples
--------
.. sourcecode:: python

    import coremltools as ct

    model = ct.models.MLModel("my_model.mlpackage")
    randomized_mlmodel = ct.models.utils.randomize_weights(mlmodel)

T)r  r  )r  _WeightRandomizer)r  randomized_mlmodels     r+   randomize_weightsrI  +  s"    2 +-/ r-   
output_dirmerge_chunks_to_pipelinecheck_output_correctnessc                    S n[        U [        [        R                  R                  45      (       d  [        S[        U 5       S35      eU" U 5      u  pVUR                  n[        U5      u  pn
UR                  S   nUR                  U   n[        R                  " SUR                   SUR                   SU S[        UR                  5       35        [        R                  " S	U	S
 S35        [        R                  " SX-
  S
 S35        [!        XX5      n[#        U" U 5      S   U5      n[        R                  " S5        [%        USSU[        R&                  R(                  U(       + S9nA[*        R,                  " 5         [        R                  " S5        [%        USSU[        R&                  R(                  U(       + S9nA[*        R,                  " 5         [        R                  " S5        U(       ao  [        R                  " S5        [        U [        5      (       a7  [        R                  R	                  U [        R&                  R(                  S9nOU n[/        UUUS9  [0        R2                  " USS9  [        U [        5      (       aG  [0        R4                  R7                  U 5      n[0        R4                  R9                  U5      u  nnUS-  nOSnU(       a  [;        UU5      n[0        R4                  R=                  UUS-   5      nUR?                  U5        U(       aW  [        R                  " S5        [        R                  R	                  U[        R&                  R(                  S9n[/        WUS9  g"g"[0        R4                  R=                  UUS-   5      n[0        R4                  R=                  UUS-   5      nUR?                  U5        UR?                  U5        [        R                  " S U S!35        g")#a  
Utility function to split a mlpackage model into two mlpackages of approximately same file size.

Parameters
----------
model: str or MLModel
    Path to the mlpackage file, or a Core ML model, to be split into two mlpackages of approximately same file size.

output_dir: str
    Path to output directory where the two model chunks / pipeline model would be saved.

    If the `model` is `{path}/{model_name}.mlpackage`, the chunk models are going to be saved as:
    1. first chunk model: `{output_dir}/{model_name}_chunk1.mlpackage`
    2. second chunk model: `{output_dir}/{model_name}_chunk2.mlpackage`
    3. chunked pipeline model: `{output_dir}/{model_name}_chunked_pipeline.mlpackage`

    If the `model` is type of `MLModel`, the chunk models are saved as:
    1. first chunk model: `{output_dir}/chunk1.mlpackage`
    2. second chunk model: `{output_dir}/chunk2.mlpackage`
    3. chunked pipeline model: `{output_dir}/chunked_pipeline.mlpackage`

merge_chunks_to_pipeline: bool
    If True, model chunks are managed inside a single pipeline model for easier asset maintenance.

check_output_correctness: bool
    - If True, compares the outputs of original Core ML model with that of pipelined CoreML model chunks and reports PSNR in dB.
    - Enabling this feature uses more memory. Disable it if your machine runs out of memory.

Examples
--------
.. sourcecode:: python

    import coremltools as ct

    model_path = "my_model.mlpackage"
    output_dir = "./output/"

    # The following code will produce two smaller models:
    # `./output/my_model_chunk1.mlpackage` and `./output/my_model_chunk2.mlpackage`
    # It also compares the output numerical of the original Core ML model with the chunked models.
    ct.models.utils.bisect_model(
        model_path,
        output_dir,
    )

    # The following code will produce a single pipeline model `./output/my_model_chunked_pipeline.mlpackage`
    ct.models.utils.bisect_model(
        model_path,
        output_dir,
        merge_chunks_to_pipeline=True,
    )

    # You can also pass the MLModel object directly
    mlmodel = ct.models.MLModel(model_path)
    ct.models.utils.bisect_model(
        mlmodel,
        output_dir,
        merge_chunks_to_pipeline=True,
    )
c                    [        U [        5      (       a  [        U 5      n[        U 5      nOU R                  nU R
                  n[        R                  " UUR                  U5      n[        UR                  5      S:  d  SUR                  ;  a  [        S5      eUR                  S   n[        UR                  5      Ul        X14$ )Nr   r  z@'bisect_model' only support model with a single 'main' function.)r'   r   r   r  r2   r?   r%  r   r  r   r   r   r/   r   )r7   ru   r?   r  r  s        r+   "get_pymil_prog_and_spec_from_model8bisect_model.<locals>.get_pymil_prog_and_spec_from_model  s     eS!!U#D3E:K;;D++K "&&%%

 t~~"fDNN&B_``~~f%t/zr-   z,'model' must be type of [str, MLModel]. Got rN  r  zThe incision op: name=z, type=z, index=ru  zFirst chunk size = z.2fz MBzSecond chunk size = r   zConverting the two programsr  r  )r  r  r  r   r  zConversion of first chunk done.z Conversion of second chunk done.z&Verifying output correctness of chunksr   )
full_modelfirst_chunk_modelsecond_chunk_modelT)exist_okr]   r   zchunked_pipeline.mlpackagez.Verifying output correctness of pipeline model)rQ  pipeline_modelzchunk1.mlpackagezchunk2.mlpackagezSaved chunks in z8 with the suffix _chunk1.mlpackage and _chunk2.mlpackageN) r'   r   r   r   r   r   rg  r  _get_op_idx_split_locationr   r   r   infor5   op_typer   _make_first_chunk_prog_make_second_chunk_progr  r   r-  _gccollect$_verify_output_correctness_of_chunksrH   makedirsrI   basenamerO   r  re   r.  )r7   rJ  rK  rL  rO  r  ru   r  op_idxfirst_chunk_weights_sizetotal_weights_size
main_blockincision_opprog_chunk1prog_chunk2model_chunk1model_chunk2r  mlpackage_namer5   r]   rU  out_path_pipelineout_path_chunk1out_path_chunk2s                            r+   bisect_modelrm  K  sn   D4 ec3::#5#5677GU}TUVWW 4E:JD,,L <VVZ;[8F&8'J''/KLL
 !1!1 2'+:M:M9NhW]V^^_`cdndydy`z_{| LL&'?&DCHILL'(:(UVY'ZZ]^_ )6K
 **51!4K LL./"*oo..55L 	KKMLL23"*oo..55L 	KKMLL34  =>eS!!jj((coo>V>V(WGG,*+	
 LLd+%**51((##N3a&|\BHHMM*d=Y6YZ-. $LLIJ ZZ//!1I1I 0 N 1"- $ ((--
D;M4MN((--
D;M4MN/*/*zl*bc	
r-   rQ  rR  rS  rU  c                   ^ SSK Jm  SSK Jn  S[        R                  S[        R                  S[
        4U4S jjn0 nU R                  R                  R                   H  nU" U5      XgR                  '   M     U R                  U5      nUb  UR                  U5      n	U	n
OUb  Ub  UR                  U5      n0 nUR                  R                  R                   HJ  nUR                  U;   a  UUR                     XR                  '   M0  XgR                     XR                  '   ML     UR                  U5      nUn
O[        S	5      eUR                  5        H  nU" X   X   U S
9  M     g)zYVerifies the end-to-end output correctness of full (original) model versus chunked modelsr   )compute_snr_and_psnr)random_gen_input_feature_typeoriginal_outputsfinal_outputs
log_prefixc           
        > SnT" X 5      u  pET" X5      u  pFXe-
  n[         R                  " U SUS SUS SUS S35        Xc:  a  [         R                  " US S35        U$ [         R                  " US SU S	35        U$ )
zReport PSNR values across two compatible tensors.
This util is from https://github.com/apple/ml-stable-diffusion/blob/main/python_coreml_stable_diffusion/torch2coreml.py#L80,
with a slightly modification.
#   z: PSNR changed by z.1fz dB (z -> rD   z dB is low!z dB > z) dB (minimum allowed) parity check passed)r   rW  warning)	rq  rr  rs  ABSOLUTE_MIN_PSNRr]   original_psnr
final_psnr	dB_changero  s	           r+   report_correctness@_verify_output_correctness_of_chunks.<locals>.report_correctness  s    
 /0@S,-=M.	l,YsO5s@SSWXbcfWgghi	
 )OOz#.k:;
  LLc"&):(;;de r-   NzFEither a single Pipeline model or two model chunks should be provided.)rq  rr  rs  )(coremltools.converters.mil.testing_utilsro  rp  r   r"  r   r2   r3   r4   r5   r   r   r0   )rQ  rR  rS  rU  rp  r{  r6   
input_descoutputs_from_full_modeloutputs_from_pipeline_modelrr  outputs_from_first_chunk_modelsecond_chunk_input_dictoutputs_from_second_chunk_modelr  ro  s                  @r+   r]  r]    sf    NS[[  be 2 J &&2288
&CJ&O
??# 9 )00<!&4&<&<Z&H#3		&+=+I):)B)B:)N& #%,22>>DDJ"@@;YOO<'8 <Foo;V'8 E +=*D*DE\*]'7abb ,0024>'1"	
 3r-   r  c                    U R                   S   nSnUR                   H  nUR                  S:X  d  M  [        UR                  R                  [
        R                  5      (       d  MJ  UR                  R                  R                  UR                  R                  R                  -  S-  nX$-  nM     US-  nSnUR                   GH  nUR                  S:X  a{  [        UR                  R                  [
        R                  5      (       aH  UR                  R                  R                  UR                  R                  R                  -  S-  nXd-  nXe:  d  M  UR                  R                  S5      (       a  M  [        UR                  5      S:X  d  M  [        UR                  S   R                  5      S:X  d  M  UR                  R                  U5      nXvU4s  $    [        S5      e)zXFind the op that approximately bisects the graph as measure by weights size on each sider  r   consti   r   r   z/Not able to find the bisect point in the model.)r   r   rX  r'   r  r   r"  sizeitemsize
startswithr   r   	child_opsr
  r   )r  rc  total_size_in_mbr  
size_in_mb	half_sizecumulative_size_in_mbr`  s           r+   rV  rV  _  sn   'J##:: Z

CKK%H%H266::+>+>>+NJ* $ !1$I ##:: Z

CKK%H%H266::+>+>>+NJ!/! ".JJ))'22BJJ1$BJJqM++,1**004F2BBB $" F
GGr-   r  r`  c                    [        5       n[        US-   5       H  nU R                  U   nUR                  R	                  S5      (       a  M4  UR
                   HX  nUR                  b  M  UR                   H6  nU R                  R                  U5      nXq:  d  M%  UR                  U5        M8     MZ     M     [        U5      $ )Nr   r  )r1   r  r   rX  r  r   r  r  r
  r:  r/   )r  r`  boundary_varsr  r  varchild_opchild_op_idxs           r+   _get_first_chunk_outputsr    s     EM6A:a zz$$W--zz77?$'MM','7'7'='=h'G'0)--c2 %2 "  r-   r  c                     / nU Ho  nUR                   [        R                  R                  :w  a  UR	                  U5        M>  [
        R                  " USUR                  S9nUR	                  U5        Mq     U$ )Nfp32r*   r   r5   )r   r&  typesr   r@  _mbcastr5   )r  r  new_boundary_varsr  fp32_vars        r+   _add_fp32_castsr    s`    99

'$$S)xx#V#((CH$$X.  r-   c                     U R                   S   n[        X!5      n[        X#5      nUR                  R	                  5         UR                  U5        [        S   " U 5        U $ )zIBuild first chunk by declaring early outputs and removing unused subgraphr  common::dead_code_elimination)r   r  r  r   clearset_outputsr(  )r  r`  r  r  r  s        r+   rY  rY    s[    
 NN6"E,U;M (=	MM	'(23D9Kr-   c           	      t   U R                   S   n[        R                  R                  Ul        [        X!5      nUR                  U   nU   U GHC  n[        UR                  UR                  [        R                  R                  :w  a  UR                  O[        R                  R                  UR                  S9nUR                  S   UR                   UR                  S   R                  '   [#        UR                   R%                  5       5      Ul        SnUR                  [        R                  R                  :X  a.  [(        R*                  " UR                  S   SUR,                  S9nOUR                  S   nUR/                  UUUSS9  GMF     SSS5        [0        S	   " U 5        [3        5       nUR                   R5                  5        H$  u  p[7        U
R8                  5      S:  d  M   XU	'   M&     Xl        [#        UR                   R%                  5       5      Ul        U $ ! , (       d  f       N= f)
zDBuild second chunk by rebuilding a pristine MIL Program from MLModelr  )	sym_shaper   r5   r   Nr   )r*   r   	before_opT)	anchor_opold_varnew_varforce_replacer  )r   r   r   iOS16opset_versionr  r   _Placeholderr  r   r&  r  r   r  r5   r   _input_dictr#  r   function_inputsr  r  r  replace_uses_of_var_after_opr(  _OrderedDictr   r   r  )r  r`  r  r  boundary_opr  new_placeholderr  new_input_dictr;   r   s              r+   rZ  rZ    s   NN6"E****E -U;M ""6*K 
 C*))#&99

#?ciiTZZ__XXO BQAXAXYZA[Eo55a8==>$)%*;*;*B*B*D$EE!GyyDJJOO+((_%<%<Q%?vY\Y_Y_`)11!4..%
 # / 	! ! 
8 23D9 "^N!!'')q{{a !1 * '!%"3"3":":"<=EKM 
s   EH))
H7ml_model	from_type(_proto.FeatureTypes_pb2.ArrayFeatureTypeto_typefunction_namesinput_namesoutput_namesz_ct.models.model.MLModelc           	        ^^^^^ [         R                  R                  R                  [         R                  R                  R                  4nS[
        R                  R                  SS4S jnSSS[        [        R                  R                     4S jnS[        [        [              S	[        S
[        [           S[        [           4S jn	SSS[        [           S[         R                  R                  S[         4S jn
S[        S[        ["        [        [        R$                  R&                     [        [        R(                     4      4UU4S jjnS[        S[        ["        [        [        R$                  R&                     [        [        R(                     4      4UU4S jjnSSS[        R(                  S[        R*                  SS4UUU4S jjnSSS[        R(                  SS4UU4S jjnU" U S9nX;  a  [-        SU 35      eTU;  a  [-        SU 35      eUT:X  a  [/        U 5      $ U" US9mU" TS9mU	" XOR0                  R2                  / S9nU	" X_R0                  R4                  S/S9n[6        R8                  " UUR:                  U R<                  S9mT(       d  S/m[>        R@                  " S T 35        T H"  nUTRB                  ;  d  M  [-        S!U S"35      e   UR0                  R2                   H@  nU
" UXAS#9(       d  M  U" URD                  S$9 H  u  nnU" UUURF                  S%   S&9  M     MB     UR0                  R4                   H2  nU
" UXQS#9(       d  M  U" URD                  S$9 H  u  nnU" UUS'9  M     M4     [I        S( TRB                  RK                  5        5       5      n[M        TS)S*UU RN                  UR0                  S+9$ ),a  
Change the tensor data types of Core ML model inputs / outputs. Supported types are FLOAT16, FLOAT32.

Parameters
----------
ml_model: MLModel
    A Core ML model that needs to change its input/output type.
    Note:
    - the original model is not modified, the model with updated types is returned as a new instance.
    - only an mlProgram is supported (not pipelines, not neural networks).

from_type:
    The type that should be changed from.

to_type:
    The type that will be used instead of all the `from_type` type.

function_names:
    Optional list of function names where the input/output needs to be changed. If not specified, only the "main"
    function will be updated.

input_names:
    Optional list of input names that should be updated (by default none of the inputs will be updated).

output_names:
    Optional list of output names that should be updated (by default all the outputs that match the `from_type`
    type will be updated).

Examples
--------
.. sourcecode:: python

    from coremltools.models.model import MLModel
    from coremltools.utils import change_input_output_tensor_type
    from coremltools.proto.FeatureTypes_pb2 import ArrayFeatureType

    model = MLModel("my_model.mlpackage")
    updated_model = change_input_output_tensor_type(
        ml_model=model,
        from_type=ArrayFeatureType.FLOAT32,
        to_type=ArrayFeatureType.FLOAT16,
    )
    updated_model.save("my_updated_model.mlpackage")
r7   rA   rp  c                     [        U [        R                  R                  5      (       d  [	        S[        U 5       S35      eU R                  5       nUR                  S5      nUS:w  a  [	        SU 35      eU$ )Nz>input model must be of type ct.models.MLModel, actual type is rD   r   rj   z7input model must be an mlProgram, actual model type is )r'   r   r   r   r   rg  r   rs   )r7   r   r  s      r+   _get_model_spec8change_input_output_tensor_type.<locals>._get_model_spec	  sv    %!3!344]^bch^i]jjklmm^^%
**62
$VWaVbcddr-   feature_typer  c                 (   U [         R                  R                  R                  :X  a  [        R
                  R                  $ U [         R                  R                  R                  :X  a  [        R
                  R                  $ [        SU  S35      e)Nzinvalid feature type: z!, supported only FLOAT16, FLOAT32)
r   FeatureTypes_pb2ri  FLOAT16r&  r  r   rk  r  r   r  s    r+   
_get_dtype3change_input_output_tensor_type.<locals>._get_dtype(	  sj     622CCKKK::??"622CCKKK::??"1,?`abbr-   names	desc_listdefaultc                     U c  Un U Vs/ s H'  nSU ;   d  UR                   U ;   d  M  UR                   PM)     sn$ s  snf )N*r5   )r  r  r  r*   s       r+   _sanitize_names8change_input_output_tensor_type.<locals>._sanitize_names1	  s:    =E )M	1SE\QVVu_	MMMs   ;;feature_descz#_proto.Model_pb2.FeatureDescription	data_typec                    U R                   U;  a'  [        R                  " SU R                    SU 35        gU R                  R	                  S5      nUS:w  a(  [        R                  " SU R                    SU S35        gU R                  R
                  R                  nXB:w  a(  [        R                  " S	U R                    S
U S35        gg)Nzignoring feature z+ as it's not in the list of required names Fr   rf  zignoring output z (type: rD   zignoring output tensor z (data type: T)r5   r   debugrg  rs   rf  rh  )r  r  r  r  feature_data_types        r+   _eligible_feature_desc?change_input_output_tensor_type.<locals>._eligible_feature_desc6	  s    
 E)MM-l.?.?-@@klqkrst#((33F;++MM,\->->,?x~UVWX(--<<EE)MM3L4E4E3FmTeSffghir-   var_namec              3     >^ #    T H  nTR                   U   n[        [        U 4S jUR                  R	                  5        5       5      S 5      nU(       d  MR  UR
                  [        R                  R                  :  aC  [        R                  " SUR                   S35        [        R                  R                  Ul        X#4v   M     g 7f)Nc              3   <   >#    U  H  u  pUT:X  d  M  Uv   M     g 7fr&   r.  )r  r;   r   r  s      r+   r  Kchange_input_output_tensor_type.<locals>._get_input_vars.<locals>.<genexpr>N	  s     O*=$!hAA*=s   	%upgrading opset_version for function 	 to iOS16)r   nextiterr   r   r  r   r   r  r   rv  r5   r  r5   r  r  r  r  s   `   r+   _get_input_vars8change_input_output_tensor_type.<locals>._get_input_varsK	  s     "D>>$'DtO$++*;*;*=OOQUVCs%%

(8(88OO&KDII;V_$`a),)9)9D&i #s   ACA5Cc              3     >^ #    T H  nTR                   U   n[        [        U 4S jUR                   5       5      S 5      nU(       d  MD  UR                  [
        R                  R                  :  aC  [        R                  " SUR                   S35        [
        R                  R                  Ul        X#4v   M     g 7f)Nc              3   J   >#    U  H  oR                   T:X  d  M  Uv   M     g 7fr&   r  )r  r   r  s     r+   r  Lchange_input_output_tensor_type.<locals>._get_output_vars.<locals>.<genexpr>X	  s     J|!vv7IAA|s   #	#r  r  )r   r  r  r   r  r   r   r  r   rv  r5   r  s   `   r+   _get_output_vars9change_input_output_tensor_type.<locals>._get_output_varsU	  s     "D>>$'DtJt||JJDQCs%%

(8(88OO&KDII;V_$`a),)9)9D&i #s   AC A5C feature_varfirst_operationNc                   > UR                      STR                  5        3nU R                  SU 3-   n[        R                  " XXBS9nUR
                  R                   R                  UR
                  UUS9  TU R                  R                  l	        [        R                  R                  TUR                  R                  5       5      Ul        S S S 5        g ! , (       d  f       g = f)Nfp_to_)r*   r   r5   r  r  r  r  )enclosing_blockget_bitwidthr5   r  r  r  r  rg  rf  rh  r&  r  tensorsym_type	get_shape	_sym_type)	r  r  r  from_dtype_strr  r*   
from_dtypeto_dtyper  s	         r+   _cast_input_type9change_input_output_tensor_type.<locals>._cast_input_type_	  s    
 ,,!*"9"9";!<=N#((T.1A+BBH;8gADD  ==$$# > 
 9@L,,5$(JJ$5$5h@T@T@^@^@`$aK! -,,s   CC
C'c                   > UR                   R                     STR                  5        3nU R                  SU 3-   n[        R
                  " XUS9nUR                   R                  R                  UR                   UUS9  X4l        X0l        TU R                  R                  l	        S S S 5        g ! , (       d  f       g = f)Nr  r  r  r  )
r  r  r  r5   r  r  r  rg  rf  rh  )r  r  to_dtype_strr  r*   r  r  s        r+   _cast_output_type:change_input_output_tensor_type.<locals>._cast_output_typep	  s     ^^++ 5 5 789L#((T,+@@H;JADD  ==$$# > 
 F (8?L,,5 ,++s   BB22
C )r7   z8not supported from_type: must be an ArrayFeatureType of z6not supported to_type: must be an ArrayFeatureType of r  )r  r  r  r  r  r  zfunctions: z
function 'z' not defined in the model)r  r  r  )r  r   )r  r  r  )r  r  c              3   L   #    U  H  oR                   R                  v   M     g 7fr&   )r  r  )r  r  s     r+   r  2change_input_output_tensor_type.<locals>.<genexpr>	  s     cKbx44::Kbs   "$r  r  )r  r  r  r   r  )(r   r  ri  r  rk  r   r   r   _Typer&  r  doubler  _Listr   r  r+  r  r  FunctionVar	Operationr   	_deepcopyr3   r4   r   r%  r   r  r?   r   r  r   r5   r   r   r   r  r  )r  r  r  r  r  r  SUPPORTED_TYPESr  r  r  r  r  r  r  r  ml_model_specr   
desc_inputr  	input_vardesc_output
output_varmodel_opset_versionr  r  r  s     ``                   @@@r+   change_input_output_tensor_typer    s   j 	00880088O
	szz11 	6N 	c@c	tzz  	!cNys4 N NUZ[^U_ Ndijmdn N
;Sz **;; 
	* #  )F9TZZEXEX;Y[deiemem[n;n4o*p     3  9VIdjjFYFY<Z\efjfnfn\o<o5p+q    b;bXXb b 
	b b"@;@JN((@	@ @  $(3M'STcSdeffo%QRaQbcddG""3Jw/H!?X?X?^?^hjkK"AZAZAaAalokpqL"" +@@!--D  MMK/01#	DNN*z)4NOPP $ $//55
%:[f#2JOO#LHi*)]e]p]pqr]st $M 6 %0077%;lh$4k>N>N$O Hj;JO %P 8 c4>>K`K`Kbcc"1++'33 r-   )NN)FN)r   F)r   F)TT)Fr&   )r  )r  r  )FT)NNN)r  r  r  gcr[  mathr   osrH   shutilrM   
subprocessrO  sysr   tempfilerE   warningsrp   collectionsr   r  r   r   	functoolsr   
_lru_cachetypingr   	_Callabler   r  r   r  r	   r  r
   r  r   r  r   r  r   r  numpyr   coremltoolsr   r   r   r   r  r   r   r   coremltools.converters.milr   r&  $coremltools.converters.mil.converterr   r  ,coremltools.converters.mil.frontend.milprotor   r%  coremltools.converters.mil.milr   r  r   _Program5coremltools.converters.mil.mil.passes.defs.preprocessr   r   4coremltools.converters.mil.mil.passes.defs.randomizer   rG  0coremltools.converters.mil.mil.passes.graph_passr   r  ,coremltools.converters.mil.mil.passes.helperr   _block_context_manager3coremltools.converters.mil.mil.passes.pass_pipeliner   r?  3coremltools.converters.mil.mil.passes.pass_registryr   r(  &coremltools.converters.mil.mil.programr   r  _depsr    rS   rG   rX   r  r\   rY   libmodelpackager$   rQ   scipy.sparsesparser  r,   r<   r   ra   rf   rz   r   r   r   r   r   r   r   ALLr   r   r   r   r   r  r&  r5  r8  r>  rA  rE  rI  rU  rZ  r_  rc  ro  r  r  r  r  r+  r  r  r  r  r4  rE  rI  rm  r]  rV  Blockr  r  r  rY  rZ  r  r.  r-   r+   <module>r+     s             3 & - (   (   ( "   "   T 3  ' 2 L S 9 > a e ` N  # " !  + ?  #'#':":3: C.: 		:zVC VC VI.X"# ""4 "J(V;J. $0#3#3 E?D;~ 9>AJ FJ\<~*.AH00.-6%
  	P$@R (,k*C/0knk 	k` 8<U[%U[dCOO34U[ 
U[p"!"" " 	"P 6'+!3!8!8#*!*)95G+HHI* * t_	*
 * *  (*+*Z(B9 B9Jaalla a 	a
 
aP#
!P#P#n !'	G0/G0*/U3s;K5L0L*MG0 G0 	G0
 
G0TF 1604	D
#++,D
D
 (oD
 (o	D
P :>:>6:	K
$K
 !56K
 ""67K
 23	K

 
K
\HT\\ HfS#s]6K HDDJJ  dhh " 4:: eDHHo $  
,, 
\\"2$,, 2 2 2r -1)-*.M"M9M 8M eCj)	M
 5:&M E#J'M  MCEMs   &M1 1M8