}YhB(dZddlmZmZddlmZgdZGddZeZeZ e ej e ej dZ e ej dZe ej d Ze eje eje eje ejd Ze eje eje eje ejd Ze eje ejd Ze eje eje eje ejd Ze ejdZe ejdZe ejdZe ej e ej!dZ"e ej#e ej#dZ$e ej%dZ&e ej%dZ'e ej(dZ)e ej(dZ*dS)aF PyTorch provides two global :class:`ConstraintRegistry` objects that link :class:`~torch.distributions.constraints.Constraint` objects to :class:`~torch.distributions.transforms.Transform` objects. These objects both input constraints and return transforms, but they have different guarantees on bijectivity. 1. ``biject_to(constraint)`` looks up a bijective :class:`~torch.distributions.transforms.Transform` from ``constraints.real`` to the given ``constraint``. The returned transform is guaranteed to have ``.bijective = True`` and should implement ``.log_abs_det_jacobian()``. 2. ``transform_to(constraint)`` looks up a not-necessarily bijective :class:`~torch.distributions.transforms.Transform` from ``constraints.real`` to the given ``constraint``. The returned transform is not guaranteed to implement ``.log_abs_det_jacobian()``. The ``transform_to()`` registry is useful for performing unconstrained optimization on constrained parameters of probability distributions, which are indicated by each distribution's ``.arg_constraints`` dict. These transforms often overparameterize a space in order to avoid rotation; they are thus more suitable for coordinate-wise optimization algorithms like Adam:: loc = torch.zeros(100, requires_grad=True) unconstrained = torch.zeros(100, requires_grad=True) scale = transform_to(Normal.arg_constraints["scale"])(unconstrained) loss = -Normal(loc, scale).log_prob(data).sum() The ``biject_to()`` registry is useful for Hamiltonian Monte Carlo, where samples from a probability distribution with constrained ``.support`` are propagated in an unconstrained space, and algorithms are typically rotation invariant.:: dist = Exponential(rate) unconstrained = torch.zeros(100, requires_grad=True) sample = biject_to(dist.support)(unconstrained) potential_energy = -dist.log_prob(sample).sum() .. note:: An example where ``transform_to`` and ``biject_to`` differ is ``constraints.simplex``: ``transform_to(constraints.simplex)`` returns a :class:`~torch.distributions.transforms.SoftmaxTransform` that simply exponentiates and normalizes its inputs; this is a cheap and mostly coordinate-wise operation appropriate for algorithms like SVI. In contrast, ``biject_to(constraints.simplex)`` returns a :class:`~torch.distributions.transforms.StickBreakingTransform` that bijects its input down to a one-fewer-dimensional space; this a more expensive less numerically stable transform but is needed for algorithms like HMC. The ``biject_to`` and ``transform_to`` objects can be extended by user-defined constraints and transforms using their ``.register()`` method either as a function on singleton constraints:: transform_to.register(my_constraint, my_transform) or as a decorator on parameterized constraints:: @transform_to.register(MyConstraintClass) def my_factory(constraint): assert isinstance(constraint, MyConstraintClass) return MyTransform(constraint.param1, constraint.param2) You can create your own registry by creating a new :class:`ConstraintRegistry` object. ) constraints transforms)_Number)ConstraintRegistry biject_to transform_toc0eZdZdZfdZddZdZxZS)rz5 Registry to link constraints to transforms. cVi|_tdSN) _registrysuper__init__)self __class__s y/var/www/tools.fuzzalab.pt/emblema-extractor/venv/lib/python3.11/site-packages/torch/distributions/constraint_registry.pyrzConstraintRegistry.__init__Us& Nc|fdSttjrtttrt tjst d|j<|S)a Registers a :class:`~torch.distributions.constraints.Constraint` subclass in this registry. Usage:: @my_registry.register(MyConstraintClass) def construct_transform(constraint): assert isinstance(constraint, MyConstraint) return MyTransform(constraint.arg_constraints) Args: constraint (subclass of :class:`~torch.distributions.constraints.Constraint`): A subclass of :class:`~torch.distributions.constraints.Constraint`, or a singleton object of the desired class. factory (Callable): A callable that inputs a constraint object and returns a :class:`~torch.distributions.transforms.Transform` object. Nc0|Sr )register)factory constraintrs rz-ConstraintRegistry.register..ls4==W#E#ErzLExpected constraint to be either a Constraint subclass or instance, but got ) isinstancer Constrainttype issubclass TypeErrorr rrrs`` rrzConstraintRegistry.registerYs$ ?EEEEE E j+"8 9 9 *j))J*d++ :  .4 4  k_ikk &-z"rc |jt|}n4#t$r'tdt|jddwxYw||S)ah Looks up a transform to constrained space, given a constraint object. Usage:: constraint = Normal.arg_constraints["scale"] scale = transform_to(constraint)(torch.zeros(1)) # constrained u = transform_to(constraint).inv(scale) # unconstrained Args: constraint (:class:`~torch.distributions.constraints.Constraint`): A constraint object. Returns: A :class:`~torch.distributions.transforms.Transform` object. Raises: `NotImplementedError` if no transform has been registered. zCannot transform z constraintsN)r rKeyErrorNotImplementedError__name__rs r__call__zConstraintRegistry.__call__|sx( nT*%5%56GG   %KD$4$4$=KKK  wz"""s 1Ar )r" __module__ __qualname____doc__rrr# __classcell__)rs@rrrPsf!!!!F#######rrctjSr )ridentity_transformrs r_transform_to_realr+s   ((rc^t|j}tj||jSr )rbase_constraintrIndependentTransformreinterpreted_batch_ndimsrbase_transforms r_biject_to_independentr2s/z9::N  * <  rc^t|j}tj||jSr )rr-rr.r/r0s r_transform_to_independentr4s/!*"<==N  * <  rc(tjSr )r ExpTransformr*s r_transform_to_positiver7s  " $ $$rctjtjtj|jdgS)N)rComposeTransformr6AffineTransform lower_boundr*s r_transform_to_greater_thanr=s=  &  # % %  &z'=q A A   rctjtjtj|jdgS)N)rr:r6r; upper_boundr*s r_transform_to_less_thanrAs=  &  # % %  &z'=r B B   rcdt|jto |jdk}t|jto |jdk}|r|rt jS|j}|j|jz }t jt jt j||gS)Nrr9)rr<rr@rSigmoidTransformr:r;)r lower_is_0 upper_is_1locscales r_transform_to_intervalrHs :)733S 8NRS8S :)733S 8NRS8S-j-*,,,  C  "Z%; ;E  &  $ & & (B3(N(NO  rc(tjSr )rStickBreakingTransformr*s r_biject_to_simplexrK  , . ..rc(tjSr )rSoftmaxTransformr*s r_transform_to_simplexrOs  & ( ((rc(tjSr )rLowerCholeskyTransformr*s r_transform_to_lower_choleskyrRrLrc(tjSr )rPositiveDefiniteTransformr*s r_transform_to_positive_definiterUs  / 1 11rc(tjSr )rCorrCholeskyTransformr*s r_transform_to_corr_choleskyrXs  + - --rc`tjd|jD|j|jS)Nc,g|]}t|Sr.0cs r z"_biject_to_cat.. ///!1///rr CatTransformcseqdimlengthsr*s r_biject_to_catrg s3  "//z///AS  rc`tjd|jD|j|jS)Nc,g|]}t|Sr[rr]s rr`z%_transform_to_cat..222Qa222rrbr*s r_transform_to_catrls3  "22*/222JNJDV  rcTtjd|jD|jS)Nc,g|]}t|Sr[r\r]s rr`z$_biject_to_stack..rarrStackTransformrdrer*s r_biject_to_stackrqs.  $//z///  rcTtjd|jD|jS)Nc,g|]}t|Sr[rjr]s rr`z'_transform_to_stack.."rkrror*s r_transform_to_stackrts.  $22*/222JN  rN)+r&torch.distributionsrr torch.typesr__all__rrrrrealr+ independentr2r4positive nonnegativer7 greater_thangreater_than_eqr= less_thanrAintervalhalf_open_intervalrHsimplexrKrOlower_choleskyrRpositive_definitepositive_semidefiniterU corr_choleskyrXcatrgrlstackrqrtr[rrrsAAF87777777   F#F#F#F#F#F#F#F#R   !!##  K$%% {'(()))(&%) K+,,-,{.//0/ K()) K+,, {+,, {.//%%0/-,-,*)% K,-- K/00 {/00 {233431010.- K)** {,--.-+* K()) K233 {+,, {56676-,43*)$ K'((//)(/{*++)),+) {122//32/{455 {89922:9652 K-.. {011..21/.. KO$$%$ {''('  K%&&'& {())*)r