efax 2.3.0__tar.gz → 2.4.0__tar.gz

efax-2.4.0/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2020, Neil Girdhar
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

{efax-2.3.0 → efax-2.4.0}/PKG-INFO

@@ -1,10 +1,12 @@
 Metadata-Version: 2.4
 Name: efax
-Version: 2.3.0
+Version: 2.4.0
 Summary: Exponential families for JAX
+Keywords: exponential-family,jax,machine-learning,probability,statistics
 Author: Neil Girdhar
 Author-email: Neil Girdhar <mistersheik@gmail.com>
 License-Expression: Apache-2.0
+License-File: LICENSE
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent
@@ -23,12 +25,15 @@ Requires-Dist: numpy>=2.0
 Requires-Dist: opt-einsum>=3.4
 Requires-Dist: optimistix>=0.0.9
 Requires-Dist: optype[numpy]>=0.8.0
-Requires-Dist: tjax>=1.4.6
+Requires-Dist: tjax>=1.5.2
 Requires-Dist: typing-extensions>=4.8
 Maintainer: Neil Girdhar
 Maintainer-email: Neil Girdhar <mistersheik@gmail.com>
 Requires-Python: >=3.12, <3.15
-Project-URL: source, https://github.com/NeilGirdhar/efax
+Project-URL: Changelog, https://github.com/NeilGirdhar/efax/blob/main/CHANGELOG.md
+Project-URL: Documentation, https://github.com/NeilGirdhar/efax#readme
+Project-URL: Issues, https://github.com/NeilGirdhar/efax/issues
+Project-URL: Source, https://github.com/NeilGirdhar/efax
 Description-Content-Type: text/x-rst
 
 .. role:: bash(code)
@@ -167,10 +172,7 @@ Every `NaturalParametrization` has methods:
 - `jeffreys_prior_density`, which returns the square root of the Fisher information
   determinant,
 - `characteristic_function`, which evaluates the characteristic function of the sufficient
-  statistics via analytic continuation of the log-normalizer (note: for distributions such
-  as `GammaNP` that keep one or more natural parameters real in `_complexify`, the result
-  is approximate — check `NaturalParametrization.characteristic_function_exact` on the
-  class before relying on every component), and
+  statistics via analytic continuation of the log-normalizer, and
 - `kl_divergence`, which is the KL divergence.
 
 Every `ExpectationParametrization` has methods:

{efax-2.3.0 → efax-2.4.0}/README.rst

@@ -134,10 +134,7 @@ Every `NaturalParametrization` has methods:
 - `jeffreys_prior_density`, which returns the square root of the Fisher information
   determinant,
 - `characteristic_function`, which evaluates the characteristic function of the sufficient
-  statistics via analytic continuation of the log-normalizer (note: for distributions such
-  as `GammaNP` that keep one or more natural parameters real in `_complexify`, the result
-  is approximate — check `NaturalParametrization.characteristic_function_exact` on the
-  class before relying on every component), and
+  statistics via analytic continuation of the log-normalizer, and
 - `kl_divergence`, which is the KL divergence.
 
 Every `ExpectationParametrization` has methods:

{efax-2.3.0 → efax-2.4.0}/efax/__init__.py

@@ -84,6 +84,7 @@ from ._src.parameter.parameter import distribution_parameter
 from ._src.parameter.ring import BooleanRing, ComplexField, IntegralRing, RealField, Ring
 from ._src.parameter.support import (
     CircularBoundedSupport,
+    HermitianMatrixSupport,
     ScalarSupport,
     SimplexSupport,
     SquareMatrixSupport,
@@ -154,6 +155,7 @@ __all__ = [
     "HasEntropyEP",
     "HasEntropyNP",
     "HasGeneralizedConjugatePrior",
+    "HermitianMatrixSupport",
     "IntegralRing",
     "InverseGammaEP",
     "InverseGammaNP",

efax-2.4.0/efax/_src/analytic_continuation.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from array_api_compat import array_namespace
+from tjax import JaxArray, JaxComplexArray, Shape
+from tjax.dataclasses import dataclass
+
+
+@dataclass
+class ComplexContinuation:
+    """Analytic continuation of a complex field's real and imaginary coordinates."""
+
+    real: JaxComplexArray
+    imag: JaxComplexArray
+
+    @property
+    def shape(self) -> Shape:
+        return self.real.shape
+
+    @property
+    def dtype(self) -> object:
+        xp = array_namespace(self.real, self.imag)
+        return xp.result_type(self.real, self.imag)
+
+    def __array_namespace__(self, *, api_version: str | None = None) -> object:  # noqa: PLW3201
+        return array_namespace(self.real, self.imag, api_version=api_version)
+
+    def __neg__(self) -> ComplexContinuation:
+        return ComplexContinuation(-self.real, -self.imag)
+
+    def __add__(self, other: JaxArray | ComplexContinuation) -> ComplexContinuation:
+        other_real, other_imag = complex_parts(other)
+        return ComplexContinuation(self.real + other_real, self.imag + other_imag)
+
+    def __radd__(self, other: JaxArray | ComplexContinuation) -> ComplexContinuation:
+        return self + other
+
+    def __sub__(self, other: JaxArray | ComplexContinuation) -> ComplexContinuation:
+        other_real, other_imag = complex_parts(other)
+        return ComplexContinuation(self.real - other_real, self.imag - other_imag)
+
+    def __rsub__(self, other: JaxArray | ComplexContinuation) -> ComplexContinuation:
+        other_real, other_imag = complex_parts(other)
+        return ComplexContinuation(other_real - self.real, other_imag - self.imag)
+
+    def __mul__(self, other: JaxArray | ComplexContinuation) -> ComplexContinuation:
+        other_real, other_imag = complex_parts(other)
+        return ComplexContinuation(
+            self.real * other_real - self.imag * other_imag,
+            self.real * other_imag + self.imag * other_real,
+        )
+
+    def __rmul__(self, other: JaxArray | ComplexContinuation) -> ComplexContinuation:
+        return self * other
+
+    def __truediv__(self, other: JaxArray | ComplexContinuation) -> ComplexContinuation:
+        other_real, other_imag = complex_parts(other)
+        denominator = other_real**2 + other_imag**2
+        return ComplexContinuation(
+            (self.real * other_real + self.imag * other_imag) / denominator,
+            (self.imag * other_real - self.real * other_imag) / denominator,
+        )
+
+
+def analytic_continue_parameter(eta: JaxArray, delta: JaxArray) -> JaxArray | ComplexContinuation:
+    """Continue real coordinates without treating complex storage as one coordinate."""
+    xp = array_namespace(eta, delta)
+    if xp.isdtype(eta.dtype, "complex floating"):
+        return ComplexContinuation(
+            xp.real(eta) + 1j * xp.real(delta),
+            xp.imag(eta) + 1j * xp.imag(delta),
+        )
+    return eta + 1j * delta
+
+
+def complex_parts(z: JaxArray | ComplexContinuation) -> tuple[JaxArray, JaxArray]:
+    """Return the real-coordinate pair represented by a complex value."""
+    if isinstance(z, ComplexContinuation):
+        return z.real, z.imag
+    xp = array_namespace(z)
+    return xp.real(z), xp.imag(z)
+
+
+def analytic_abs_square(z: JaxArray | ComplexContinuation) -> JaxArray:
+    """Return ``real(z)**2 + imag(z)**2`` without conjugating continued coordinates."""
+    real, imag = complex_parts(z)
+    return real**2 + imag**2
+
+
+def complex_value(z: JaxArray | ComplexContinuation) -> JaxArray:
+    """Return the ordinary complex value represented by real and imaginary coordinates."""
+    real, imag = complex_parts(z)
+    return real + 1j * imag

{efax-2.3.0 → efax-2.4.0}/efax/_src/distributions/chi_square.py

@@ -6,7 +6,7 @@ from typing import override
 import jax.random as jr
 import jax.scipy.special as jss
 from array_api_compat import array_namespace
-from tjax import JaxArray, JaxRealArray, KeyArray, Shape
+from tjax import JaxArray, JaxRealArray, KeyArray, Shape, complex_gammaln
 from tjax.dataclasses import dataclass
 
 from efax._src.interfaces.samplable import Samplable
@@ -46,7 +46,7 @@ class ChiSquareNP(
     @override
     def log_normalizer(self) -> JaxRealArray:
         k_over_two = self.k_over_two_minus_one + 1.0
-        return jss.gammaln(k_over_two) - k_over_two * math.log(0.5)
+        return complex_gammaln(k_over_two) - k_over_two * math.log(0.5)
 
     @override
     def to_exp(self) -> ChiSquareEP:

{efax-2.3.0 → efax-2.4.0}/efax/_src/distributions/cmvn/circularly_symmetric.py

@@ -5,15 +5,24 @@ from typing import override
 
 import jax.random as jr
 from array_api_compat import array_namespace
-from tjax import JaxArray, JaxComplexArray, JaxRealArray, KeyArray, Shape, outer_product
+from tjax import (
+    JaxArray,
+    JaxComplexArray,
+    JaxRealArray,
+    KeyArray,
+    Shape,
+    complex_logdet,
+    sesquilinear_outer,
+)
 from tjax.dataclasses import dataclass
 
+from efax._src.analytic_continuation import ComplexContinuation, complex_value
 from efax._src.interfaces.multidimensional import Multidimensional
 from efax._src.interfaces.samplable import Samplable
 from efax._src.mixins.has_entropy import HasEntropyEP, HasEntropyNP
 from efax._src.natural_parametrization import NaturalParametrization
 from efax._src.parameter import (
-    SymmetricMatrixSupport,
+    HermitianMatrixSupport,
     VectorSupport,
     complex_field,
     distribution_parameter,
@@ -36,7 +45,7 @@ class ComplexCircularlySymmetricNormalNP(
     """
 
     negative_precision: JaxComplexArray = distribution_parameter(
-        SymmetricMatrixSupport(hermitian=True, negative_semidefinite=True)
+        HermitianMatrixSupport(negative_semidefinite=True)
     )
     # S = -1/negative_precision, U = 0
     # P = S.conjugate, R = 0
@@ -57,8 +66,12 @@ class ComplexCircularlySymmetricNormalNP(
 
     @override
     def log_normalizer(self) -> JaxRealArray:
-        xp = array_namespace(self)
-        log_det_s = xp.log(xp.real(xp.linalg.det(-self.negative_precision)))
+        negative_precision = complex_value(self.negative_precision)
+        xp = array_namespace(self.negative_precision)
+        log_det_s = complex_logdet(-negative_precision)
+        # Ordinary parameters must keep the log-normalizer real-valued.
+        if not isinstance(self.negative_precision, ComplexContinuation):
+            log_det_s = xp.real(log_det_s)
         return -log_det_s + self.dimensions() * math.log(math.pi)
 
     @override
@@ -76,7 +89,7 @@ class ComplexCircularlySymmetricNormalNP(
     def sufficient_statistics(
         cls, x: JaxComplexArray, **fixed_parameters: JaxArray
     ) -> ComplexCircularlySymmetricNormalEP:
-        return ComplexCircularlySymmetricNormalEP(outer_product(x, x))
+        return ComplexCircularlySymmetricNormalEP(sesquilinear_outer(x, x))
 
     @override
     def sample(self, key: KeyArray, shape: Shape | None = None) -> JaxComplexArray:
@@ -100,7 +113,7 @@ class ComplexCircularlySymmetricNormalEP(
     """
 
     variance: JaxComplexArray = distribution_parameter(
-        SymmetricMatrixSupport(hermitian=True, positive_semidefinite=True)
+        HermitianMatrixSupport(positive_semidefinite=True)
     )
 
     @property

{efax-2.3.0 → efax-2.4.0}/efax/_src/distributions/cmvn/unit_variance.py

@@ -8,6 +8,7 @@ from array_api_compat import array_namespace
 from tjax import JaxArray, JaxComplexArray, JaxRealArray, KeyArray, Shape, abs_square
 from tjax.dataclasses import dataclass
 
+from efax._src.analytic_continuation import analytic_abs_square
 from efax._src.interfaces.multidimensional import Multidimensional
 from efax._src.interfaces.samplable import Samplable
 from efax._src.mixins.has_entropy import HasEntropyEP, HasEntropyNP
@@ -50,9 +51,9 @@ class ComplexMultivariateUnitVarianceNormalNP(
 
     @override
     def log_normalizer(self) -> JaxRealArray:
-        xp = array_namespace(self)
-        mean = self.two_mean * 0.5
-        return xp.sum(abs_square(mean), axis=-1) + self.dimensions() * math.log(math.pi)
+        two_mean_abs_square = analytic_abs_square(self.two_mean)
+        xp = array_namespace(two_mean_abs_square)
+        return 0.25 * xp.sum(two_mean_abs_square, axis=-1) + self.dimensions() * math.log(math.pi)
 
     @override
     def to_exp(self) -> ComplexMultivariateUnitVarianceNormalEP:

{efax-2.3.0 → efax-2.4.0}/efax/_src/distributions/complex_normal/complex_normal.py

@@ -8,6 +8,7 @@ from array_api_compat import array_namespace
 from tjax import JaxArray, JaxComplexArray, JaxRealArray, KeyArray, Shape, abs_square
 from tjax.dataclasses import dataclass
 
+from efax._src.analytic_continuation import complex_parts
 from efax._src.interfaces.samplable import Samplable
 from efax._src.mixins.has_entropy import HasEntropyEP, HasEntropyNP
 from efax._src.natural_parametrization import NaturalParametrization
@@ -62,18 +63,15 @@ class ComplexNormalNP(
 
     @override
     def log_normalizer(self) -> JaxRealArray:
-        xp = array_namespace(self)
-        pseudo_precision = xp.conj(self.pseudo_precision)
-        _, s, mu = self._r_s_mu()
-        det_s = s
-        det_h = -self.negative_precision
-        return (
-            -abs_square(mu) * self.negative_precision
-            - xp.real(xp.square(mu) * pseudo_precision)
-            + 0.5 * xp.log(det_s)
-            - 0.5 * xp.log(det_h)
-            + math.log(math.pi)
-        )
+        xp = array_namespace(self.negative_precision)
+        eta_x, eta_y = complex_parts(self.mean_times_precision)
+        pseudo_x, pseudo_y = complex_parts(self.pseudo_precision)
+        h_xx = self.negative_precision + pseudo_x
+        h_xy = pseudo_y
+        h_yy = self.negative_precision - pseudo_x
+        det_h = h_xx * h_yy - h_xy**2
+        h_inv_quadratic = (h_yy * eta_x**2 - 2.0 * h_xy * eta_x * eta_y + h_xx * eta_y**2) / det_h
+        return math.log(math.pi) - 0.5 * xp.log(det_h) - 0.25 * h_inv_quadratic
 
     @override
     def to_exp(self) -> ComplexNormalEP:
@@ -96,12 +94,12 @@ class ComplexNormalNP(
         return ComplexNormalEP(x, abs_square(x), xp.square(x))
 
     def _r_s_mu(self) -> tuple[JaxComplexArray, JaxRealArray, JaxComplexArray]:
-        xp = array_namespace(self)
+        xp = array_namespace(self.negative_precision)
         r = -xp.conj(self.pseudo_precision) / self.negative_precision
         s = xp.reciprocal((abs_square(r) - 1.0) * self.negative_precision)
         k = xp.conj(self.pseudo_precision) / self.negative_precision
         l_eta = (
-            0.5
+            xp.asarray(0.5)
             * xp.conj(self.mean_times_precision)
             / ((abs_square(k) - 1.0) * self.negative_precision)
         )

{efax-2.3.0 → efax-2.4.0}/efax/_src/distributions/complex_normal/unit_variance.py

@@ -8,6 +8,7 @@ from array_api_compat import array_namespace
 from tjax import JaxArray, JaxComplexArray, JaxRealArray, KeyArray, Shape, abs_square
 from tjax.dataclasses import dataclass
 
+from efax._src.analytic_continuation import analytic_abs_square
 from efax._src.interfaces.samplable import Samplable
 from efax._src.mixins.has_entropy import HasEntropyEP, HasEntropyNP
 from efax._src.natural_parametrization import NaturalParametrization
@@ -48,8 +49,7 @@ class ComplexUnitVarianceNormalNP(
 
     @override
     def log_normalizer(self) -> JaxRealArray:
-        mean = self.two_mean * 0.5
-        return abs_square(mean) + math.log(math.pi)
+        return 0.25 * analytic_abs_square(self.two_mean) + math.log(math.pi)
 
     @override
     def to_exp(self) -> ComplexUnitVarianceNormalEP:

{efax-2.3.0 → efax-2.4.0}/efax/_src/distributions/complex_von_mises.py

@@ -8,12 +8,15 @@ from array_api_compat import array_namespace
 from tjax import JaxArray, JaxComplexArray, JaxRealArray, Shape, inverse_softplus, softplus
 from tjax.dataclasses import dataclass
 
+from efax._src.analytic_continuation import analytic_abs_square
 from efax._src.mixins.exp_to_nat.exp_to_nat import ExpToNat
 from efax._src.mixins.has_entropy import HasEntropyEP, HasEntropyNP
 from efax._src.natural_parametrization import NaturalParametrization
 from efax._src.parameter import ComplexField, ScalarSupport, complex_field, distribution_parameter
 from efax._src.parametrization import SimpleDistribution
 
+from .von_mises import _log_bessel_iv_over_power
+
 
 @dataclass
 class ComplexVonMisesNP(
@@ -43,9 +46,9 @@ class ComplexVonMisesNP(
 
     @override
     def log_normalizer(self) -> JaxRealArray:
-        xp = array_namespace(self)
-        kappa = xp.abs(self.mean_times_concentration)
-        return kappa + math.log(2.0 * math.pi) + xp.log(jss.i0e(kappa))
+        rho = analytic_abs_square(self.mean_times_concentration)
+        xp = array_namespace(rho)
+        return math.log(2.0 * math.pi) + _log_bessel_iv_over_power(xp.asarray(0.0), rho)
 
     @override
     def to_exp(self) -> ComplexVonMisesEP:

{efax-2.3.0 → efax-2.4.0}/efax/_src/distributions/dirichlet_common.py

@@ -5,7 +5,7 @@ from typing import Any, override
 import jax.random as jr
 import jax.scipy.special as jss
 from array_api_compat import array_namespace
-from tjax import JaxRealArray, KeyArray, Shape
+from tjax import JaxRealArray, KeyArray, Shape, complex_gammaln
 from tjax.dataclasses import dataclass
 
 from efax._src.interfaces.multidimensional import Multidimensional
@@ -33,7 +33,7 @@ class DirichletCommonNP[EP: "DirichletCommonEP[Any]"](
     def log_normalizer(self) -> JaxRealArray:
         xp = array_namespace(self)
         q = self.alpha_minus_one
-        return xp.sum(jss.gammaln(q + 1.0), axis=-1) - jss.gammaln(
+        return xp.sum(complex_gammaln(q + 1.0), axis=-1) - complex_gammaln(
             xp.sum(q, axis=-1) + self.dimensions()
         )
 

{efax-2.3.0 → efax-2.4.0}/efax/_src/distributions/gamma.py

@@ -1,11 +1,19 @@
 from __future__ import annotations
 
-from typing import ClassVar, Self, override
+from typing import override
 
 import jax.random as jr
 import jax.scipy.special as jss
 from array_api_compat import array_namespace
-from tjax import JaxArray, JaxRealArray, KeyArray, Shape, inverse_softplus, softplus
+from tjax import (
+    JaxArray,
+    JaxRealArray,
+    KeyArray,
+    Shape,
+    complex_gammaln,
+    inverse_softplus,
+    softplus,
+)
 from tjax.dataclasses import dataclass
 
 from efax._src.interfaces.samplable import Samplable
@@ -39,8 +47,6 @@ class GammaNP(
         shape_minus_one: The shape minus one.
     """
 
-    characteristic_function_exact: ClassVar[bool] = False
-
     negative_rate: JaxRealArray = distribution_parameter(ScalarSupport(ring=negative_support))
     shape_minus_one: JaxRealArray = distribution_parameter(
         ScalarSupport(ring=RealField(minimum=-1.0))
@@ -60,7 +66,7 @@ class GammaNP(
     def log_normalizer(self) -> JaxRealArray:
         xp = array_namespace(self)
         shape = self.shape_minus_one + 1.0
-        return jss.gammaln(shape) - shape * xp.log(-self.negative_rate)
+        return complex_gammaln(shape) - shape * xp.log(-self.negative_rate)
 
     @override
     def to_exp(self) -> GammaEP:
@@ -81,17 +87,6 @@ class GammaNP(
         xp = array_namespace(x)
         return GammaEP(x, xp.log(x))
 
-    @override
-    def _complexify(self, t: Self) -> Self:  # ty: ignore
-        """Only complexify negative_rate; jss.gammaln does not accept complex inputs.
-
-        shape_minus_one is kept real, so ``characteristic_function`` cannot
-        query the log-x sufficient statistic: that component returns 1 instead
-        of the true E[exp(i·t·log x)].  Consequently,
-        ``characteristic_function_exact = False`` on this class.
-        """
-        return type(self)(self.negative_rate + 1j * t.negative_rate, self.shape_minus_one)
-
     @override
     def sample(self, key: KeyArray, shape: Shape | None = None) -> JaxRealArray:
         xp = array_namespace(self)