o ?e@s&dZddlZddlZddlZddlmZddlmZddlmZddlm Z ddlm Z ddl m Z dd l m Z dd l mZdd l mZdd l mZdd l mZddl mZddl mZddlmZ     dWddZddZ dXddZddZdYddZ      dZddZddZd d!Zd"d#Zd$d%Z d&d'Z!d(d)Z" *d[d+d*Z# , -d\d.d/Z$d]d1d0Z%d^d2d3Z&d_d5d4Z'd`d7d6Z( 8dad9d8Z)d:d;Z*dd?Z,d@dAZ-dbdBdCZ.dbdDdEZ/dcdFdGZ0     dddHdIZ1dYdJdKZ2dLdMZ3 dYdNdOZ4dedQdRZ5dSdTZ6GdUdVdVZ7dS)fz(Utilities for probability distributions.N) constant_op)dtypes)ops) tensor_shape) tensor_util) array_ops)array_ops_stack) check_ops)cond)control_flow_ops) linalg_ops)math_ops)nn) tf_inspectassert_integer_formc Cstj|||gdbtj|dd}|jjr tWdS|p&d|}|durPztj tj tj tj tj tji|jj}WntyOtd|jjwtj|tt|||j||||dWdS1snwYdS)axAssert that x has integer components (or floats equal to integers). Args: x: Floating-point `Tensor` data: The tensors to print out if the condition is `False`. Defaults to error message and first few entries of `x` and `y`. summarize: Print this many entries of each tensor. message: A string to prefix to the default message. int_dtype: A `tf.dtype` used to cast the float to. The default (`None`) implies the smallest possible signed int will be used for casting. name: A name for this operation (optional). Returns: Op raising `InvalidArgumentError` if `cast(x, int_dtype) != x`. valuesxnameNz{} has non-integer componentszUnrecognized type {})data summarizemessager)r name_scopeconvert_to_tensordtype is_integerr Zno_opformatrfloat16int16float32int32float64int64 base_dtypeKeyError TypeErrorrr assert_equalr cast)rrrr int_dtyperr*i/home/www/facesmatcher.com/pyenv/lib/python3.10/site-packages/tensorflow/python/ops/distributions/util.pyr&s4 $cCs t|}tt||g|SN)rZmatrix_transposer with_dependenciesr r')matrixZmatrix_tr*r*r+assert_symmetricRs r/$embed_check_nonnegative_integer_formcCstj||gd0tj|dd}tj|d|dg}|jjs+|t|d|dg7}t ||WdS1s;wYdS)z>Assert x is a non-negative tensor, and optionally of integers.rrrz'{}' must be non-negative.rz*'{}' cannot contain fractional components.N) rrrr assert_non_negativerrrrr r-)rr assertionsr*r*r+r0Xs  $csPtjddtjddfdd}tttt|ddS)zReturns whether a and b have the same dynamic shape. Args: a: `Tensor` b: `Tensor` Returns: `bool` `Tensor` representing if both tensors have the same shape. arbc sBtttttgdtttgdS)Nr)r Z reduce_allequalrconcatshaper*r4r5r*r+all_shapes_equalzsz,same_dynamic_shape..all_shapes_equalcSs tdS)NF)rconstantr*r*r*r+s z$same_dynamic_shape..)rrtf_condr r r6rrank)r4r5r:r*r9r+same_dynamic_shapejs  r?cCsR|dur|Szt|}Wn ty|}Ynw|dus!|dur#|St||S)a`Helper which tries to return a static value. Given `x`, extract it's value statically, optionally casting to a specific dtype. If this is not possible, None is returned. Args: x: `Tensor` for which to extract a value statically. dtype: Optional dtype to cast to. Returns: Statically inferred value if possible, otherwise None. N)rconstant_valuer&nparray)rrZx_r*r*r+maybe_get_static_values   rCFget_logits_and_probsc Cstj|||gd|du|dukrtd|durRtj|d|d}|jjs*td|rB|r2t|}|tj |ddfWdS|t j |ddfWdStj|d|d}|jjsbtd |rtd <t d |j}t|g}|rt|}|tjt |d |d dg7}n |tj||ddg7}t||}Wdn1swYtd4|rt ||fWdWdSt |t d||fWdWdS1swYWddS1swYdS)a Converts logit to probabilities (or vice-versa), and returns both. Args: logits: Floating-point `Tensor` representing log-odds. probs: Floating-point `Tensor` representing probabilities. multidimensional: Python `bool`, default `False`. If `True`, represents whether the last dimension of `logits` or `probs`, a `[N1, N2, ... k]` dimensional tensor, representing the logit or probability of `shape[-1]` classes. validate_args: Python `bool`, default `False`. When `True`, either assert `0 <= probs <= 1` (if not `multidimensional`) or that the last dimension of `probs` sums to one. name: A name for this operation (optional). dtype: `tf.DType` to prefer when converting args to `Tensor`s. Returns: logits, probs: Tuple of `Tensor`s. If `probs` has an entry that is `0` or `1`, then the corresponding entry in the returned logit will be `-Inf` and `Inf` respectively. Raises: ValueError: if neither `probs` nor `logits` were passed in, or both were. rNz(Must pass probs or logits, but not both.logitsrrz!logits must having floating type.probsrz probs must having floating type.Zvalidate_probsg?zprobs does not sum to 1.r1z$probs has components greater than 1.g)rr ValueErrorrr is_floatingr&#embed_check_categorical_event_shaperZsoftmaxr Zsigmoidrr;r r2Z assert_near reduce_sumassert_less_equalr r-loglog1p)rErGZmultidimensional validate_argsrrone dependenciesr*r*r+rDs^     1'"cCs tjdtjdtjdi|jdS)z7Helper returning True if dtype is known to be unsigned.TF)rboolZuint8Zuint16getr$dtr*r*r+_is_known_unsigned_by_dtypes  rWcCs8tjdtjdtjdtjdtjdtjdtjdi|j dS)z5Helper returning True if dtype is known to be signed.TF) rrr r"Zint8rr!r#rTr$rUr*r*r+_is_known_signed_by_dtypes rXcCst|pt|S)z(Helper returning True if dtype is known.)rWrXrUr*r*r+_is_known_dtype srYcCspt|s td|j|jrtdt|jj dS|j r&t |jj S|j tjkr0tdStd|j)zDHelper returning the largest integer exactly representable by dtype.Unrecognized dtype: {})rYr&rrrJintrAfinfoas_numpy_dtypeZnmantrZiinfomaxr$rrSrUr*r*r+_largest_integer_by_dtypes racCs0t|s td|jt|rdSdt|S)zEHelper returning the smallest integer exactly representable by dtype.rZrrH)rYr&rrrWrarUr*r*r+_smallest_integer_by_dtypes  rbcCs*t|s td|j|jp|jtjkS)z7Helper returning True if dtype.is_integer or is `bool`.rZ)rYr&rrrr$rrSrUr*r*r+_is_integer_like_by_dtype%srcrKc CsVtj||gdtj|dd}|jj}|jrt|nd}|dkr)td|j z | d}Wn t y=t dwt |dd url|jdj}|d krUt d ||krct d |j |||Wd Stj|d dd}ttj|dddtjt|dd d dtj||d|j |dg|Wd S1swYd S)aEmbeds checks that categorical distributions don't have too many classes. A categorical-type distribution is one which, e.g., returns the class label rather than a one-hot encoding. E.g., `Categorical(probs)`. Since distributions output samples in the same dtype as the parameters, we must ensure that casting doesn't lose precision. That is, the `parameter.dtype` implies a maximum number of classes. However, since shape is `int32` and categorical variables are presumed to be indexes into a `Tensor`, we must also ensure that the number of classes is no larger than the largest possible `int32` index, i.e., `2**31-1`. In other words the number of classes, `K`, must satisfy the following condition: ```python K <= min( int(2**31 - 1), # Largest float as an index. { dtypes.float16: int(2**11), # Largest int as a float16. dtypes.float32: int(2**24), dtypes.float64: int(2**53), }.get(categorical_param.dtype.base_dtype, 0)) ``` Args: categorical_param: Floating-point `Tensor` representing parameters of distribution over categories. The rightmost shape is presumed to be the number of categories. name: A name for this operation (optional). Returns: categorical_param: Input `Tensor` with appropriate assertions embedded. Raises: TypeError: if `categorical_param` has an unknown `dtype`. ValueError: if we can statically identify `categorical_param` as being too large (for being closed under int32/float casting). rcategorical_paramrrz3Unable to validate size of unrecognized dtype ({}).r\zDA categorical-distribution parameter must have at least 1 dimension.rHNr[zAA categorical-distribution parameter must have at least 2 events.zZNumber of classes exceeds `dtype` precision, i.e., {} implies shape ({}) cannot exceed {}.Zx_shaper1zSNumber of classes exceeds `dtype` precision, i.e., {} dtype cannot exceed {} shape.)rrrrr$rJrar&rr get_shapewith_rank_at_leastrIrdimension_valuedimsvaluerr8r r-r Zassert_rank_at_leastassert_greater_equalrM)rdrrZx_dtypeZmax_event_sizeZx_shape_staticZ event_sizer*r*r+rK,sd)   ! $Tembed_check_casting_closedc Cstj||gdtj|dd}t|js"|jjs"td|jjt|s1|js1td|jt|jsFt|sFtd||jj|jg}|rT|t j |ddg7}|jjrg|t ||d |jd g7}n8t |jt |kr|t j |t |d t |dg7}|st|jt|kr|t j|t|d t|dg7}|s|Wd St||Wd S1swYd S)aEnsures integers remain unaffected despite casting to/from int/float types. Example integer-types: `uint8`, `int32`, `bool`. Example floating-types: `float32`, `float64`. The largest possible integer representable by an IEEE754 floating-point is `2**(1 + mantissa_bits)` yet the largest possible integer as an int-type is `2**(bits - 1) - 1`. This function ensures that a `Tensor` purporting to have integer-form values can be cast to some other type without loss of precision. The smallest representable integer is the negative of the largest representable integer, except for types: `uint8`, `uint16`, `bool`. For these types, the smallest representable integer is `0`. Args: x: `Tensor` representing integer-form values. target_dtype: TF `dtype` under which `x` should have identical values. assert_nonnegative: `bool` indicating `x` should contain nonnegative values. name: A name for this operation (optional). Returns: x: Input `Tensor` with appropriate assertions embedded. Raises: TypeError: if `x` is neither integer- nor floating-type. TypeError: if `target_dtype` is neither integer- nor floating-type. TypeError: if neither `x` nor `target_dtype` are integer-type. rrrz+{}.dtype must be floating- or integer-type.z4target_dtype ({}) must be floating- or integer-type.zIAt least one of {}.dtype ({}) and target_dtype ({}) must be integer-type.zElements must be non-negative.r1zElements must be {}-equivalent.)r)rzElements cannot exceed {}.z#Elements cannot be smaller than {}.N)rrrrcrrJr&rrr r2rrarMrbrjr r-)rZ target_dtypeZassert_nonnegativerr3r*r*r+"embed_check_integer_casting_closeds!    7$rllog_combinationscCstj|||gd0tj|dd}tj|dd}t|d}t|d}tj|dgd}||WdS1s= 2`, where the last two dimensions are equal. transform: Element-wise function mapping `Tensors` to `Tensors`. To be applied to the diagonal of `matrix`. If `None`, `matrix` is returned unchanged. Defaults to `None`. name: A name to give created ops. Defaults to "matrix_diag_transform". Returns: A `Tensor` with same shape and `dtype` as `matrix`. matrix_diag_transformr.rN)rrrrZmatrix_diag_partZmatrix_set_diag)r.Z transformrdiagZtransformed_diagZtransformed_matr*r*r+rss2    rsrotate_transposec Csdtj|||gdtj|dd}tj|dd}t|t|}|j}|durk|durk|dkr<|WdSt |t ||}|dkrT|WdSt t ||}tj||dWdSt|}tt|dt| ||t||}td|}t||}t||gd}tj||dWdS1swYdS) a)Circularly moves dims left or right. Effectively identical to: ```python numpy.transpose(x, numpy.roll(numpy.arange(len(x.shape)), shift)) ``` When `validate_args=False` additional graph-runtime checks are performed. These checks entail moving data from to GPU to CPU. Example: ```python x = tf.random.normal([1, 2, 3, 4]) # Tensor of shape [1, 2, 3, 4]. rotate_transpose(x, -1).shape == [2, 3, 4, 1] rotate_transpose(x, -2).shape == [3, 4, 1, 2] rotate_transpose(x, 1).shape == [4, 1, 2, 3] rotate_transpose(x, 2).shape == [3, 4, 1, 2] rotate_transpose(x, 7).shape == rotate_transpose(x, 3).shape # [2, 3, 4, 1] rotate_transpose(x, -7).shape == rotate_transpose(x, -3).shape # [4, 1, 2, 3] ``` Args: x: `Tensor`. shift: `Tensor`. Number of dimensions to transpose left (shift<0) or transpose right (shift>0). name: Python `str`. The name to give this op. Returns: rotated_x: Input `Tensor` with dimensions circularly rotated by shift. Raises: TypeError: if shift is not integer type. rrrshiftNr[r)perm)rrrr Zassert_integerrr@rendimsrAsignabsZrollZarangerZ transposer>where_v2r lessmodranger7)rrvrZshift_value_staticrxrwfirstlastr*r*r+ruFs<$           $ pick_vectorc Cstj||||fdvtj|dd}|jtjkr#td||jtjft|}|dur9|r0|n|WdStj|dd}tj|dd}|j|jkrYtd||j||jft |d }t t ||gd t |d |gt ||d gWdS1swYdS) aPicks possibly different length row `Tensor`s based on condition. Value `Tensor`s should have exactly one dimension. If `cond` is a python Boolean or `tf.constant` then either `true_vector` or `false_vector` is immediately returned. I.e., no graph nodes are created and no validation happens. Args: cond: `Tensor`. Must have `dtype=tf.bool` and be scalar. true_vector: `Tensor` of one dimension. Returned when cond is `True`. false_vector: `Tensor` of one dimension. Returned when cond is `False`. name: Python `str`. The name to give this op. Example: ```python pick_vector(tf.less(0, 5), tf.range(10, 12), tf.range(15, 18)) # [10, 11] pick_vector(tf.less(5, 0), tf.range(10, 12), tf.range(15, 18)) # [15, 16, 17] ``` Returns: true_or_false_vector: `Tensor`. Raises: TypeError: if `cond.dtype != tf.bool` TypeError: if `cond` is not a constant and `true_vector.dtype != false_vector.dtype` rr rz%s.dtype=%s which is not %sN true_vector false_vectorz&%s.dtype=%s does not match %s.dtype=%srrH)rrrrrrSr&rr@rr8slicer7r{where)r rrrZcond_value_staticrnr*r*r+rs0     $prefer_static_broadcast_shapecstj|||gdCddfdd}fdd}||}||}|dur7|dur7t||WdS||}||}t||WdS1sOwYdS) aConvenience function which statically broadcasts shape when possible. Args: shape1: `1-D` integer `Tensor`. Already converted to tensor! shape2: `1-D` integer `Tensor`. Already converted to tensor! name: A string name to prepend to created ops. Returns: The broadcast shape, either as `TensorShape` (if broadcast can be done statically), or as a `Tensor`. rcSstj|dtjdS)Nr8rF)rrrr!rr*r*r+make_shape_tensorsz8prefer_static_broadcast_shape..make_shape_tensorcs4t|tjr|St|}|durt|SdSr,) isinstancer TensorShaperr@)sZs_rr*r+get_tensor_shapes  z7prefer_static_broadcast_shape..get_tensor_shapecs0t|tjs |S|r|Std)Nz6Cannot broadcast from partially defined `TensorShape`.)rrris_fully_definedas_listrI)rrr*r+get_shape_tensors  z7prefer_static_broadcast_shape..get_shape_tensorN)rrrZbroadcast_static_shapeZbroadcast_dynamic_shape)Zshape1Zshape2rrrZshape1_Zshape2_r*rr+rs    $cCtt|S)zReturn static rank of tensor `x` if available, else `tf.rank(x)`. Args: x: `Tensor` (already converted). Returns: Numpy array (if static rank is obtainable), else `Tensor`. )prefer_static_valuerr>rr*r*r+prefer_static_rank rcCr)zReturn static shape of tensor `x` if available, else `tf.shape(x)`. Args: x: `Tensor` (already converted). Returns: Numpy array (if static shape is obtainable), else `Tensor`. )rrr8rr*r*r+prefer_static_shaperrcCst|}|dur |S|S)zReturn static value of tensor `x` if available, else `x`. Args: x: `Tensor` (already converted). Returns: Numpy array (if static value is obtainable), else `Tensor`. N)rr@)rZstatic_xr*r*r+rs rcCs>|durdSt||d}tt|dddd@S)z2Generate a new seed, from the given seed and salt.Nzutf-8i)strencoder]hashlibmd5 hexdigest)seedsaltstringr*r*r+ gen_new_seeds rc Cstj|d|gdtj|dd}t|jdddurSt|jj dj }t dd |d }|t |krAt d |t|}|jdd||g}n+t|d}tjt dtjd |tjd tjd }|jdddddg}t|}|r|tj|d|df|dgdg}n|d|dftj||dgdg}|r|ntjt|dd||ggdd}ttj|dd|}tj||rdnd|rdndd}|||WdS1swYdS)aCreates a (batch of) triangular matrix from a vector of inputs. Created matrix can be lower- or upper-triangular. (It is more efficient to create the matrix as upper or lower, rather than transpose.) Triangular matrix elements are filled in a clockwise spiral. See example, below. If `x.get_shape()` is `[b1, b2, ..., bB, d]` then the output shape is `[b1, b2, ..., bB, n, n]` where `n` is such that `d = n(n+1)/2`, i.e., `n = int(np.sqrt(0.25 + 2. * m) - 0.5)`. Example: ```python fill_triangular([1, 2, 3, 4, 5, 6]) # ==> [[4, 0, 0], # [6, 5, 0], # [3, 2, 1]] fill_triangular([1, 2, 3, 4, 5, 6], upper=True) # ==> [[1, 2, 3], # [0, 5, 6], # [0, 0, 4]] ``` For comparison, a pure numpy version of this function can be found in `util_test.py`, function `_fill_triangular`. Args: x: `Tensor` representing lower (or upper) triangular elements. upper: Python `bool` representing whether output matrix should be upper triangular (`True`) or lower triangular (`False`, default). name: Python `str`. The name to give this op. Returns: tril: `Tensor` with lower (or upper) triangular elements filled from `x`. Raises: ValueError: if `x` cannot be mapped to a triangular matrix. fill_triangularrrrr\rHNg?@g?zGInput right-most shape ({}) does not correspond to a triangular matrix.r[r.rpr)Z num_lowerZ num_upper)rrrrrgr8rfrAr!rhrisqrtfloorrIr concatenaterr r(rr rreverserrr7reshapeZmatrix_band_part set_shape) rupperrmrnstatic_final_shaperxZx_listZ new_shaper*r*r+rsJ+ %&$" $rc Cstj|d|gdtj|dd}t|jdddur?t|jj dj }t||dd}|jdd  |g}nt |d}||dd}|jddd  dg}t |}|rw|d d ddf}|d ddddf}nt j|d dddf|dgd }|d ddddf}t jt j||dgd |dgd } || } t | t jt |dd ||dggd d } t j|| d d||fgdd } | || WdS1swYdS) aCreates a vector from a (batch of) triangular matrix. The vector is created from the lower-triangular or upper-triangular portion depending on the value of the parameter `upper`. If `x.shape` is `[b1, b2, ..., bB, n, n]` then the output shape is `[b1, b2, ..., bB, d]` where `d = n (n + 1) / 2`. Example: ```python fill_triangular_inverse( [[4, 0, 0], [6, 5, 0], [3, 2, 1]]) # ==> [1, 2, 3, 4, 5, 6] fill_triangular_inverse( [[1, 2, 3], [0, 5, 6], [0, 0, 4]], upper=True) # ==> [1, 2, 3, 4, 5, 6] ``` Args: x: `Tensor` representing lower (or upper) triangular elements. upper: Python `bool` representing whether output matrix should be upper triangular (`True`) or lower triangular (`False`, default). name: Python `str`. The name to give this op. Returns: flat_tril: (Batch of) vector-shaped `Tensor` representing vectorized lower (or upper) triangular elements from `x`. fill_triangular_inverserrrr[rHNr\.rrp)rrrrrgr8rfrAr!rhrirrrrrr7r) rrrrnrrrxZinitial_elementsZtriangular_portionZrotated_triangular_portionZconsolidated_matrixZ end_sequenceyr*r*r+rsD&"(" $rcCsdd}dd}t|d|||gT|dur.tj|dd}t||d dd d df}|dur>tj|d d}t|}|durYtj|d d}t||d d ddd f}||||WdS1siwYdS)a'Creates a matrix with values set above, below, and on the diagonal. Example: ```python tridiag(below=[1., 2., 3.], diag=[4., 5., 6., 7.], above=[8., 9., 10.]) # ==> array([[ 4., 8., 0., 0.], # [ 1., 5., 9., 0.], # [ 0., 2., 6., 10.], # [ 0., 0., 3., 7.]], dtype=float32) ``` Warning: This Op is intended for convenience, not efficiency. Args: below: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the below diagonal part. `None` is logically equivalent to `below = 0`. diag: `Tensor` of shape `[B1, ..., Bb, d]` corresponding to the diagonal part. `None` is logically equivalent to `diag = 0`. above: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the above diagonal part. `None` is logically equivalent to `above = 0`. name: Python `str`. The name to give this op. Returns: tridiag: `Tensor` with values set above, below and on the diagonal. Raises: ValueError: if all inputs are `None`. cSsFtjt|dddggdd}tj||jd}tj|||gddS)zBPrepends and appends a zero to every vector in a batch of vectors.NrHr\rrpr)rr7r8Zzerosr)rr8zr*r*r+_pads"ztridiag.._padcWsBd}|D]}|dur q|dur|}q||7}q|durtd|S)z&Adds list of Tensors, ignoring `None`.Nz6Must specify at least one of `below`, `diag`, `above`.)rI)rrrr*r*r+_adds ztridiag.._addtridiagNbelowr.rHr\rtabove)rrrrZ matrix_diag)rrtrrrrr*r*r+rs!    $rc CsTt|d||gtj|dd}|dur8tj|||d}|r/t|}||fWdS|WdStj||jdd}|tt |}tj ||dd} t t | t | | } t|t|| } tj| ||d} |s|t| |} t| }| t|| }|r||fWdS|WdS1swYdS) aComputes `log(abs(sum(weight * exp(elements across tensor dimensions))))`. If all weights `w` are known to be positive, it is more efficient to directly use `reduce_logsumexp`, i.e., `tf.reduce_logsumexp(logx + tf.math.log(w))` is more efficient than `du.reduce_weighted_logsumexp(logx, w)`. Reduces `input_tensor` along the dimensions given in `axis`. Unless `keep_dims` is true, the rank of the tensor is reduced by 1 for each entry in `axis`. If `keep_dims` is true, the reduced dimensions are retained with length 1. If `axis` has no entries, all dimensions are reduced, and a tensor with a single element is returned. This function is more numerically stable than log(sum(w * exp(input))). It avoids overflows caused by taking the exp of large inputs and underflows caused by taking the log of small inputs. For example: ```python x = tf.constant([[0., 0, 0], [0, 0, 0]]) w = tf.constant([[-1., 1, 1], [1, 1, 1]]) du.reduce_weighted_logsumexp(x, w) # ==> log(-1*1 + 1*1 + 1*1 + 1*1 + 1*1 + 1*1) = log(4) du.reduce_weighted_logsumexp(x, w, axis=0) # ==> [log(-1+1), log(1+1), log(1+1)] du.reduce_weighted_logsumexp(x, w, axis=1) # ==> [log(-1+1+1), log(1+1+1)] du.reduce_weighted_logsumexp(x, w, axis=1, keep_dims=True) # ==> [[log(-1+1+1)], [log(1+1+1)]] du.reduce_weighted_logsumexp(x, w, axis=[0, 1]) # ==> log(-1+5) ``` Args: logx: The tensor to reduce. Should have numeric type. w: The weight tensor. Should have numeric type identical to `logx`. axis: The dimensions to reduce. If `None` (the default), reduces all dimensions. Must be in the range `[-rank(input_tensor), rank(input_tensor))`. keep_dims: If true, retains reduced dimensions with length 1. return_sign: If `True`, returns the sign of the result. name: A name for the operation (optional). Returns: lswe: The `log(abs(sum(weight * exp(x))))` reduced tensor. sign: (Optional) The sign of `sum(weight * exp(x))`. reduce_weighted_logsumexplogxrN)rqkeepdimswrrT)rrrr Zreduce_logsumexpr ones_likerrNrzZ reduce_maxr{Zis_infZ zeros_likeryexprLZsqueeze) rrrqZ keep_dimsZ return_signrZlsweZsgnZ log_absw_xZmax_log_absw_xZwx_over_max_absw_xZsum_wx_over_max_absw_xr*r*r+rs>@   $rc Cstj|d|gdYtj|dd}tt|jjjd}t |t |}t || }t |}|}t t ||t ||}|t t |  }t ||t |||WdS1sewYdS)acComputes the inverse softplus, i.e., x = softplus_inverse(softplus(x)). Mathematically this op is equivalent to: ```none softplus_inverse = log(exp(x) - 1.) ``` Args: x: `Tensor`. Non-negative (not enforced), floating-point. name: A name for the operation (optional). Returns: `Tensor`. Has the same type/shape as input `x`. softplus_inverserrrrN)rrrrArNr^rr_Zepsr r|rZgreaterrr{ logical_orrexpm1)rr thresholdZ is_too_smallZ is_too_largeZtoo_small_valueZtoo_large_valuerr*r*r+r{s"  $rcCs6t|jt||}|dur|St||S)z)Returns the size of a specific dimension.N)rrgr8rfrArzr)rrqrr*r*r+dimension_sizes rc Cst|d|g|durItjjjdd\}}||j}||j}|tjj |ddd}tj |d|d }tj |d |d }||fWdSt |\}}tj |d|d }tj |d |d }|t j |dd dd d }dd}||||}}|dur|dur||krt d||n2|rtjt|d dt|d dddg} t| t|}t|}Wdn1swY||fWdS1swYdS)aValidates quadrature grid, probs or computes them as necessary. Args: quadrature_grid_and_probs: Python pair of `float`-like `Tensor`s representing the sample points and the corresponding (possibly normalized) weight. When `None`, defaults to: `np.polynomial.hermite.hermgauss(deg=8)`. dtype: The expected `dtype` of `grid` and `probs`. validate_args: Python `bool`, default `False`. When `True` distribution parameters are checked for validity despite possibly degrading runtime performance. When `False` invalid inputs may silently render incorrect outputs. name: Python `str` name prefixed to Ops created by this class. Returns: quadrature_grid_and_probs: Python pair of `float`-like `Tensor`s representing the sample points and the corresponding (possibly normalized) weight. Raises: ValueError: if `quadrature_grid_and_probs is not None` and `len(quadrature_grid_and_probs[0]) != len(quadrature_grid_and_probs[1])` !process_quadrature_grid_and_probsNr)degr\T)ordrgridrFrGZunnormalized_probsrH)rrqrrcSst|jddS)z:Returns the static size of a specific dimension or `None`.r\rH)rrgr8rfrr*r*r+_static_event_sizesz=process_quadrature_grid_and_probs.._static_event_sizezm`quadrature_grid_and_probs` must be a `tuple` of same-length zero-th-dimension `Tensor`s (saw lengths {}, {})rpzX`quadrature_grid_and_probs` must be a `tuple` of same-length zero-th-dimension `Tensor`sr1)rrrAZ polynomialZhermiteZ hermgaussZastyper_ZlinalgZnormrtupler rIrr r'rZcontrol_dependenciesridentity) Zquadrature_grid_and_probsrrPrrrGrrrnr3r*r*r+rsJ        $rr\c Cst|d|||gtj|dd}tj||jdd}tj|dd}|jjs/td|jj|s7|s7td|j j d urA|j j nt j |d d}tj|d d}t |}|d ur|}|d krb||}t |} |d ksq|j j d ur|j d |} t| d urd n t|j || ||} |j |d d } | | | } qd } n t |d k|||}d } t j|t jt|r|nd|r|ndg|d |tjd|d}| d ur|| |Wd S1swYd S)aZPads `value` to the front and/or back of a `Tensor` dim, `count` times. Args: x: `Tensor` input. axis: Scalar `int`-like `Tensor` representing the single dimension to pad. (Negative indexing is supported.) front: Python `bool`; if `True` the beginning of the `axis` dimension is padded with `value`, `count` times. If `False` no front padding is made. back: Python `bool`; if `True` the end of the `axis` dimension is padded with `value`, `count` times. If `False` no end padding is made. value: Scalar `int`-like `Tensor` representing the actual value added to the front and/or back of the `axis` dimension of `x`. count: Scalar `int`-like `Tensor` representing number of elements added to the front and/or back of the `axis` dimension of `x`. E.g., if `front = back = True` then `2 * count` elements are added. name: Python `str` name prefixed to Ops created by this function. Returns: pad: The padded version of input `x`. Raises: ValueError: if both `front` and `back` are `False`. TypeError: if `count` is not `int`-like. padrrrircountz(`count.dtype` (`{}`) must be `int`-like.z/At least one of `front`, `back` must be `True`.Nrxrqrr\rH)indicesdepthrqZon_valuer)ZpaddingsZconstant_values)rrrrrr&rrrIr8rxrr>rr@rrZdimension_at_indexrr{rZone_hotrstackrr!r)rrqZfrontbackrirrrxZaxis_Zcount_headmiddletailZ final_shaper*r*r+rsd   $rcCsdtjtjdd\}}}}||i||i}i}|D] }||||<q!|||S)aLReturns parent frame arguments. When called inside a function, returns a dictionary with the caller's function arguments. These are positional arguments and keyword arguments (**kwargs), while variable arguments (*varargs) are excluded. When called at global scope, this will return an empty dictionary, since there are no arguments. WARNING: If caller function argument names are overloaded before invoking this method, then values will reflect the overloaded value. For this reason, we recommend calling `parent_frame_arguments` at the beginning of the function. r\r)r_inspect getargvaluesrpopupdate) arg_namesZvariable_arg_nameZkeyword_arg_name local_varsZ keyword_argsZ final_argsZarg_namer*r*r+parent_frame_argumentsIs   rc@s"eZdZdZdddZddZdS) AppendDocstringaLHelper class to promote private subclass docstring to public counterpart. Example: ```python class TransformedDistribution(Distribution): @distribution_util.AppendDocstring( additional_note="A special note!", kwargs_dict={"foo": "An extra arg."}) def _prob(self, y, foo=None): pass ``` In this case, the `AppendDocstring` decorator appends the `additional_note` to the docstring of `prob` (not `_prob`) and adds a new `kwargs` section with each dictionary item as a bullet-point. For a more detailed example, see `TransformedDistribution`. NcCs||_|rHg}t|D],}||}tdd|Dr"td||}d|vr0td||d||fq |jdd|7_dSdS) aInitializes the AppendDocstring object. Args: additional_note: Python string added as additional docstring to public version of function. kwargs_dict: Python string/string dictionary representing specific kwargs expanded from the **kwargs input. Raises: ValueError: if kwargs_dict.key contains whitespace. ValueError: if kwargs_dict.value contains newlines. css|]}|VqdSr,)isspace).0rr*r*r+ sz+AppendDocstring.__init__..z(Parameter name "%s" contains whitespace. z1Parameter description for "%s" contains newlines.z * `%s`: %sz ##### `kwargs`: N)_additional_notesortedkeysanyrIlstripappendjoin)selfZadditional_noteZ kwargs_dictZbulletskeyrir*r*r+__init__s  zAppendDocstring.__init__csDtfdd}|jdur|j|_|S|jd|j7_|S)Ncs|i|Sr,r*)argskwargsfnr*r+_fnsz%AppendDocstring.__call__.._fnz %s) functoolswraps__doc__r)rrrr*rr+__call__s zAppendDocstring.__call__)rN)__name__ __module__ __qualname__rrrr*r*r*r+rns  r)NNNNr)r0r,)NNFFrDN)rK)Trk)rm)NN)ru)r)r)FN)NNNN)NNFFN)FFrr\N)8rrrnumpyrAZtensorflow.python.frameworkrrrrrZtensorflow.python.opsrrr r r=r r r rZtensorflow.python.utilrrr/r0r?rCrDrWrXrYrarbrcrKrlrmrsrurrrrrrrrrrrrrrrrr*r*r*r+s               ,   Q    a [ " > M0 -    s EE b9 CG%