U
    ‰d¸  ã                   @   s4   d dl mZ d dlmZ d dlZG dd„ deƒZdS )é    )Úschema)Ú
ModelLayerNc                       s&   e Zd Zd‡ fdd„	Zdd„ Z‡  ZS )ÚConstantWeightNÚconstant_weightc                    sœ   t tˆ ƒj|||f|Ž t tjˆ  d¡¡ˆ _ˆ j	 
¡ ˆ _tˆ jƒˆ _ˆd k	rRˆn‡ fdd„tˆ jƒD ƒ‰tˆƒˆ jks|t‚‡ ‡fdd„tˆ jƒD ƒˆ _d S )Nr   c                    s   g | ]}d ˆ j  ‘qS )g      ð?)Únum)Ú.0Ú_)Úself© úH/tmp/pip-unpacked-wheel-ua33x9lu/caffe2/python/layers/constant_weight.pyÚ
<listcomp>   s     z+ConstantWeight.__init__.<locals>.<listcomp>c                    s,   g | ]$}ˆ j  d ˆ j|f tˆ| ƒ¡‘qS )z%s_weight_%d)ÚmodelZadd_global_constantÚnameÚfloat)r   Úi©r	   Úweightsr
   r   r   "   s
   þ 
ÿ)Úsuperr   Ú__init__r   ZScalarÚnpZfloat32Zget_next_blob_referenceÚoutput_schemaÚinput_recordZfield_blobsÚdataÚlenr   ÚrangeÚAssertionErrorr   )r	   r   r   r   r   Úkwargs©Ú	__class__r   r   r      s,    ÿ  ÿÿ ÿþýzConstantWeight.__init__c                 C   s(   |  dd„ t| j| jƒD ƒ|  ¡ ¡ d S )Nc                 S   s   g | ]}|D ]}|‘qqS r
   r
   )r   Zx_w_pairÚbr
   r
   r   r   *   s       z*ConstantWeight.add_ops.<locals>.<listcomp>)ZWeightedSumÚzipr   r   r   )r	   Únetr
   r
   r   Úadd_ops(   s    þzConstantWeight.add_ops)Nr   )Ú__name__Ú
__module__Ú__qualname__r   r"   Ú__classcell__r
   r
   r   r   r      s     ûr   )Zcaffe2.pythonr   Zcaffe2.python.layers.layersr   Znumpyr   r   r
   r
   r
   r   Ú<module>   s   