U
    dV                     @   s    d dl Z G dd de jdZdS )    Nc                       sf   e Zd Z fddZeejdd Zedd Zej	dd Zejdd	 Z
d
d Zdd Z  ZS )SamplingTrainableMixinc                    s"   t t| j|| d | _d| _d S )NF)superr   __init___train_param_blobs_train_param_blobs_frozen)selfargskwargs	__class__ Q/tmp/pip-unpacked-wheel-ua33x9lu/caffe2/python/layers/sampling_trainable_mixin.pyr      s    zSamplingTrainableMixin.__init__c                 C   s   dS )z<
        List of parameter blobs for prediction net
        Nr   r   r   r   r   param_blobs   s    z"SamplingTrainableMixin.param_blobsc                 C   s   | j dkr| j| _| j S )zU
        If train_param_blobs is not set before used, default to param_blobs
        N)r   r   train_param_blobsr   r   r   r   r      s    
z(SamplingTrainableMixin.train_param_blobsc                 C   s&   | j r
t|d k	std| _ || _d S )NT)r   AssertionErrorr   )r   Zblobsr   r   r   r   #   s    
c                 C   s   dS )zG
        Add ops to the given net, using the given param_blobs
        Nr   )r   netr   r   r   r   _add_ops*   s    zSamplingTrainableMixin._add_opsc                 C   s   |  || j d S N)r   r   r   r   r   r   r   add_ops1   s    zSamplingTrainableMixin.add_opsc                 C   s   |  || j d S r   )r   r   r   r   r   r   add_train_ops4   s    z$SamplingTrainableMixin.add_train_ops)__name__
__module____qualname__r   propertyabcabstractmethodr   r   setterr   r   r   __classcell__r   r   r
   r   r      s   


r   )	metaclass)r   ABCMetar   r   r   r   r   <module>   s   