o Ei_$@sddlZddlZddlZddlmZddlmZddl m Z m Z m Z ej dZddZGdddZdd d ZGd d d ejZGd ddejZGdddejZdS)N)partial)OptionalCallableAnyjaxcCsR|t|jvrt|j|durt|j|}|S||j||}t|j|<|SN)varsmodulerng)concretenameshapedtypeinittensorrL/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/einx/nn/equinox.pycreate_or_retrieve s rc@sLeZdZGdddejjjZGdddejjjZGdddejj Z dS) ParamFactoryc@eZdZddZddZdS)zParamFactory.ConcretecCs@||_||_||_|durt|dr|j}nd}||_||_dS)Nrfloat32)r r rhasattrrr )selfr r rrr rrr__init__s  zParamFactory.Concrete.__init__cCs|t|j|j|jfSr)rCacheKeyr rrrrrrto_value_and_key%sz&ParamFactory.Concrete.to_value_and_keyN)__name__ __module__ __qualname__rrrrrrConcretes r c@s,eZdZddZddZddZddZd S) zParamFactory.CacheKeycC||_||_||_dSrr rrrr rrrrrr) zParamFactory.CacheKey.__init__cCst|j|j|jfSr)hashr rrrrrr__hash__.szParamFactory.CacheKey.__hash__cCs0t|tjo|j|jko|j|jko|j|jkSr) isinstancerrr rr)rotherrrr__eq__1s    zParamFactory.CacheKey.__eq__cCst|j|j|j}||fSr)rTracerr rr)rbackend virtual_argxrrr to_tracer9szParamFactory.CacheKey.to_tracerN)rrrrr&r)r.rrrrr(s  rc@r)zParamFactory.TracercCr!rr"r#rrrr>r$zParamFactory.Tracer.__init__cCsT|jdur|jn|dd}|jdur|jn|dd}|jdur$|jn|dd}|dur2td|dur:tdt|tr|dksG|dkrPtjj j dd }nH|d kr^tjj j d |d }n:|d krltjj j d|d }n,|dkrtjj |d|d|d}ntd|dt|t tfrtjj j ||d }tjjt|||||gtj|dS)Nr rrz/Must specify name for tensor factory eqx.Modulez/Must specify init for tensor factory eqx.Moduleget_at rearrangeg{Gz?)stddevaddg)rmultiplyg?dotin_axisout_axis batch_axisz3Don't know which initializer to use for operation '')argsoutput)r getrr ValueErrorr'strtjaxnn initializersnormalconstant lecun_normalintfloateinxtracerapplyrTensor)rr kwargsr rrrrr__call__Cs4   zParamFactory.Tracer.__call__N)rrrrrKrrrrr*=s r*N) rrrrFrGinputInputr r TensorFactoryr*rrrrrsrcCst|||||S)aCreate a tensor factory for Equinox parameters. Args: module: The module to create the parameter in. Must be an instance of ``eqx.Module``. name: Name of the parameter. If ``None``, uses a default name determined from the calling operation. Defaults to ``None``. init: Initializer for the parameter. If ``None``, uses a default init method determined from the calling operation. Defaults to ``None``. dtype: Data type of the parameter. If ``None``, uses the ``dtype`` member of the calling module or ``float32`` if it does not exist. Defaults to ``None``. Returns: A tensor factory with the given default parameters. )rr )r r rrr rrrparamdsrOc@seZdZUdZeed<eed<eed<eed<eed<eed<eej ed<eej ed <ee ed <e ed <eed <eed <e ed<         ddededededed ed ee d e d ed e de fddZ dddZdS)NormaNormalization layer. Args: stats: Einstein string determining the axes along which mean and variance are computed. Will be passed to ``einx.reduce``. params: Einstein string determining the axes along which learnable parameters are applied. Will be passed to ``einx.elementwise``. Defaults to ``"b... [c]"``. mean: Whether to apply mean normalization. Defaults to ``True``. var: Whether to apply variance normalization. Defaults to ``True``. scale: Whether to apply a learnable scale according to ``params``. Defaults to ``True``. bias: Whether to apply a learnable bias according to ``params``. Defaults to ``True``. epsilon: A small float added to the variance to avoid division by zero. Defaults to ``1e-5``. fastvar: Whether to use a fast variance computation. Defaults to ``True``. dtype: Data type of the weights. Defaults to ``"float32"``. decay_rate: Decay rate for exponential moving average of mean and variance. If ``None``, no moving average is applied. Defaults to ``None``. **kwargs: Additional parameters that specify values for single axes, e.g. ``a=4``. statsparamsmeanvar use_scaleuse_biasscalebias decay_rateepsilonfastvarrrJb... [c]TNh㈵>rc Ksb|durtd||_||_||_||_||_||_d|_d|_||_ ||_ | |_ | |_ | |_ dS)Nz6Stateful layers are currently not supported in Equinox)r<rQrRrSrTrUrVrWrXrYrZr[rrJ) rrQrRrSrTrWrXrYrZr[rrJrrrrs z Norm.__init__c Csftjj||j|jf|j|j|jrt|d|dnd|j r"t|d|dnd|j |j d|j \}}}|S)NrWr r rX)rSrTrWrXrZr[) rFr?normrQrRrSrTrUrOrVrZr[rJ)rr-r _mean_varrrrrKs  z Norm.__call__) r\TTTTNr]Trr)rrr__doc__r=__annotations__boolrrArrayrEdictrrrKrrrrrPvs`      rPc @sheZdZUdZeed<ejed<eejed<e ed<e ed<ddede d e de fd d Z dd dZ d S)LinearaLinear layer. Args: expr: Einstein string determining the axes along which the weight matrix is multiplied. Will be passed to ``einx.dot``. bias: Whether to apply a learnable bias. Defaults to ``True``. dtype: Data type of the weights. Defaults to ``"float32"``. **kwargs: Additional parameters that specify values for single axes, e.g. ``a=4``. exprweightrXrVrJTrrcKs"||_||_d|_d|_||_dSr)rhrVrirXrJ)rrhrXrrJrrrrs  zLinear.__init__NcCsBtjj||jf|jdurt|d|dndt|d|dd|jS)NrXr^ri)rXri)rFr?linearrhrVrOrJrr-r rrrrKs zLinear.__call__)Trr)rrrrbr=rcrrerrdrfrrrKrrrrrgs  rgc @sVeZdZUdZeed<eed<eed<eed<d dededede fddZ d d Z d S) DropoutaDropout layer. Args: expr: Einstein string determining the axes along which dropout is applied. Will be passed to ``einx.elementwise``. drop_rate: Drop rate. inference: Whether the layer is used in inference mode (i.e. not apply dropout). Defaults to ``False``. **kwargs: Additional parameters that specify values for single axes, e.g. ``a=4``. rh drop_raterJ inferenceFcKs||_||_||_||_dSr)rhrmrJrn)rrhrmrnrJrrrrs zDropout.__init__cCs,|jstjj||jf|j|d|jS|S)N)rmr )rnrFr?dropoutrhrmrJrkrrrrKszDropout.__call__N)F) rrrrbr=rcrErfrdrrrKrrrrrls   rl)NNNN)rFrequinoxeqx functoolsr jax.numpynumpyjnptypingrrrrGimport_r>rrrOModulerPrgrlrrrrs    OQ"