Add SymmetricMatrixSpace and ConstantDiagonalSymmetricMatrixSpace.

    - These can be used for symmetric matrices and any open subset of them (e.g. positive definite matrices) since they will share the same tangent space.
    - Also refactor spaces / tests to their own files.

PiperOrigin-RevId: 529547439

Author: srvasude

tensorflow_probability/python/distributions/BUILD

@@ -314,7 +314,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:sigmoid",
-        "//tensorflow_probability/python/experimental/tangent_spaces",
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:batched_rejection_sampler",
         "//tensorflow_probability/python/internal:distribution_util",
@@ -494,6 +493,7 @@ multi_substrate_py_library(
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:softmax_centered",
         "//tensorflow_probability/python/bijectors:softplus",
+        "//tensorflow_probability/python/experimental/tangent_spaces:simplex",
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
@@ -1631,7 +1631,6 @@ multi_substrate_py_library(
         ":distribution",
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:softmax_centered",
-        "//tensorflow_probability/python/experimental/tangent_spaces",
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
@@ -1651,7 +1650,6 @@ multi_substrate_py_library(
         ":gamma",
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:sigmoid",
-        "//tensorflow_probability/python/experimental/tangent_spaces",
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
@@ -1926,6 +1924,7 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/bijectors:shift",
         "//tensorflow_probability/python/bijectors:sigmoid",
         "//tensorflow_probability/python/bijectors:softplus",
+        "//tensorflow_probability/python/experimental/tangent_spaces:spherical",
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
@@ -2005,6 +2004,7 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/bijectors:exp",
         "//tensorflow_probability/python/bijectors:softmax_centered",
         "//tensorflow_probability/python/bijectors:softplus",
+        "//tensorflow_probability/python/experimental/tangent_spaces:simplex",
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
@@ -2360,6 +2360,7 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/bijectors:sigmoid",
         "//tensorflow_probability/python/bijectors:softmax_centered",
         "//tensorflow_probability/python/bijectors:square",
+        "//tensorflow_probability/python/experimental/tangent_spaces:spherical",
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:reparameterization",
@@ -2450,6 +2451,7 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/bijectors:softmax_centered",
         "//tensorflow_probability/python/bijectors:softplus",
         "//tensorflow_probability/python/bijectors:square",
+        "//tensorflow_probability/python/experimental/tangent_spaces:spherical",
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",

tensorflow_probability/python/distributions/dirichlet.py

@@ -316,8 +316,8 @@ def _default_event_space_bijector(self):
         validate_args=self.validate_args)
 
   def _experimental_tangent_space(self, x):
-    from tensorflow_probability.python.experimental.tangent_spaces import spaces  # pylint:disable=g-import-not-at-top
-    return spaces.ProbabilitySimplexSpace()
+    from tensorflow_probability.python.experimental.tangent_spaces import simplex  # pylint:disable=g-import-not-at-top
+    return simplex.ProbabilitySimplexSpace()
 
   def _sample_control_dependencies(self, x):
     """Checks the validity of a sample."""

tensorflow_probability/python/distributions/power_spherical.py

@@ -352,8 +352,8 @@ def _default_event_space_bijector(self):
     ], validate_args=self.validate_args)
 
   def _experimental_tangent_space(self, x):
-    from tensorflow_probability.python.experimental.tangent_spaces import spaces  # pylint:disable=g-import-not-at-top
-    return spaces.SphericalSpace()
+    from tensorflow_probability.python.experimental.tangent_spaces import spherical  # pylint:disable=g-import-not-at-top
+    return spherical.SphericalSpace()
 
   def _parameter_control_dependencies(self, is_init):
     if not self.validate_args:

tensorflow_probability/python/distributions/relaxed_onehot_categorical.py

@@ -563,6 +563,6 @@ def _default_event_space_bijector(self):
         validate_args=self.validate_args)
 
   def _experimental_tangent_space(self, x):
-    from tensorflow_probability.python.experimental.tangent_spaces import spaces  # pylint:disable=g-import-not-at-top
-    return spaces.ProbabilitySimplexSpace()
+    from tensorflow_probability.python.experimental.tangent_spaces import simplex  # pylint:disable=g-import-not-at-top
+    return simplex.ProbabilitySimplexSpace()
 

tensorflow_probability/python/distributions/spherical_uniform.py

@@ -199,8 +199,8 @@ def _default_event_space_bijector(self):
     ], validate_args=self.validate_args)
 
   def _experimental_tangent_space(self, x):
-    from tensorflow_probability.python.experimental.tangent_spaces import spaces  # pylint:disable=g-import-not-at-top
-    return spaces.SphericalSpace()
+    from tensorflow_probability.python.experimental.tangent_spaces import spherical  # pylint:disable=g-import-not-at-top
+    return spherical.SphericalSpace()
 
   def _sample_control_dependencies(self, samples):
     inner_sample_dim = samples.shape[-1]

tensorflow_probability/python/distributions/von_mises_fisher.py

@@ -495,8 +495,8 @@ def _default_event_space_bijector(self):
     ], validate_args=self.validate_args)
 
   def _experimental_tangent_space(self, x):
-    from tensorflow_probability.python.experimental.tangent_spaces import spaces  # pylint:disable=g-import-not-at-top
-    return spaces.SphericalSpace()
+    from tensorflow_probability.python.experimental.tangent_spaces import spherical  # pylint:disable=g-import-not-at-top
+    return spherical.SphericalSpace()
 
   def _parameter_control_dependencies(self, is_init):
     if not self.validate_args:

tensorflow_probability/python/experimental/tangent_spaces/BUILD

@@ -32,7 +32,25 @@ multi_substrate_py_library(
     name = "tangent_spaces",
     srcs = ["__init__.py"],
     deps = [
+        ":simplex",
         ":spaces",
+        ":spherical",
+        ":symmetric_matrix",
+    ],
+)
+
+multi_substrate_py_library(
+    name = "spaces_test_util",
+    testonly = True,
+    srcs = ["spaces_test_util.py"],
+    deps = [
+        ":spaces",
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/bijectors:identity",
+        "//tensorflow_probability/python/internal:tensorshape_util",
+        "//tensorflow_probability/python/internal:test_util",
+        "//tensorflow_probability/python/math:gradient",
     ],
 )
 
@@ -42,9 +60,7 @@ multi_substrate_py_library(
     deps = [
         # numpy dep,
         # tensorflow dep,
-        "//tensorflow_probability/python/experimental/linalg:linear_operator_row_block",
         "//tensorflow_probability/python/internal:distribution_util",
-        "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:nest_util",
         "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -56,16 +72,16 @@ multi_substrate_py_test(
     name = "spaces_test",
     size = "medium",
     srcs = ["spaces_test.py"],
-    shard_count = 10,
+    shard_count = 3,
     tags = [
         "tf1-broken",
     ],
     deps = [
         ":spaces",
+        ":spaces_test_util",
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:exp",
-        "//tensorflow_probability/python/bijectors:identity",
         "//tensorflow_probability/python/bijectors:reshape",
         "//tensorflow_probability/python/bijectors:scale",
         "//tensorflow_probability/python/bijectors:scale_matvec_tril",
@@ -74,3 +90,106 @@ multi_substrate_py_test(
         "//tensorflow_probability/python/internal:test_util",
     ],
 )
+
+multi_substrate_py_library(
+    name = "simplex",
+    srcs = ["simplex.py"],
+    deps = [
+        ":spaces",
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/experimental/linalg:linear_operator_row_block",
+        "//tensorflow_probability/python/internal:distribution_util",
+        "//tensorflow_probability/python/internal:dtype_util",
+        "//tensorflow_probability/python/internal:nest_util",
+        "//tensorflow_probability/python/internal:prefer_static",
+    ],
+)
+
+multi_substrate_py_test(
+    name = "simplex_test",
+    size = "medium",
+    srcs = ["simplex_test.py"],
+    shard_count = 3,
+    tags = [
+        "tf1-broken",
+    ],
+    deps = [
+        ":simplex",
+        ":spaces_test_util",
+        # numpy dep,
+        # scipy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/bijectors:exp",
+        "//tensorflow_probability/python/bijectors:identity",
+        "//tensorflow_probability/python/bijectors:scale",
+        "//tensorflow_probability/python/bijectors:scale_matvec_tril",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)
+
+multi_substrate_py_library(
+    name = "spherical",
+    srcs = ["spherical.py"],
+    deps = [
+        ":spaces",
+        # tensorflow dep,
+        "//tensorflow_probability/python/experimental/linalg:linear_operator_row_block",
+        "//tensorflow_probability/python/internal:distribution_util",
+        "//tensorflow_probability/python/internal:prefer_static",
+    ],
+)
+
+multi_substrate_py_test(
+    name = "spherical_test",
+    size = "medium",
+    srcs = ["spherical_test.py"],
+    shard_count = 3,
+    tags = [
+        "tf1-broken",
+    ],
+    deps = [
+        ":spaces_test_util",
+        ":spherical",
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/bijectors:exp",
+        "//tensorflow_probability/python/bijectors:scale",
+        "//tensorflow_probability/python/bijectors:scale_matvec_tril",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)
+
+multi_substrate_py_library(
+    name = "symmetric_matrix",
+    srcs = ["symmetric_matrix.py"],
+    deps = [
+        ":spaces",
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/internal:dtype_util",
+        "//tensorflow_probability/python/internal:prefer_static",
+        "//tensorflow_probability/python/math:linalg",
+    ],
+)
+
+multi_substrate_py_test(
+    name = "symmetric_matrix_test",
+    size = "medium",
+    srcs = ["symmetric_matrix_test.py"],
+    shard_count = 3,
+    tags = [
+        "tf1-broken",
+    ],
+    deps = [
+        ":spaces_test_util",
+        ":symmetric_matrix",
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/bijectors:exp",
+        "//tensorflow_probability/python/bijectors:reshape",
+        "//tensorflow_probability/python/bijectors:scale",
+        "//tensorflow_probability/python/internal:test_util",
+        "//tensorflow_probability/python/math:linalg",
+    ],
+)

tensorflow_probability/python/experimental/tangent_spaces/__init__.py

@@ -14,21 +14,25 @@
 # ============================================================================
 """TensorFlow Probability experimental tangent spaces package."""
 
+from tensorflow_probability.python.experimental.tangent_spaces.simplex import ProbabilitySimplexSpace
 from tensorflow_probability.python.experimental.tangent_spaces.spaces import AxisAlignedSpace
 from tensorflow_probability.python.experimental.tangent_spaces.spaces import FullSpace
 from tensorflow_probability.python.experimental.tangent_spaces.spaces import GeneralSpace
-from tensorflow_probability.python.experimental.tangent_spaces.spaces import ProbabilitySimplexSpace
-from tensorflow_probability.python.experimental.tangent_spaces.spaces import SphericalSpace
 from tensorflow_probability.python.experimental.tangent_spaces.spaces import TangentSpace
 from tensorflow_probability.python.experimental.tangent_spaces.spaces import UnspecifiedTangentSpaceError
 from tensorflow_probability.python.experimental.tangent_spaces.spaces import ZeroSpace
+from tensorflow_probability.python.experimental.tangent_spaces.spherical import SphericalSpace
+from tensorflow_probability.python.experimental.tangent_spaces.symmetric_matrix import ConstantDiagonalSymmetricMatrixSpace
+from tensorflow_probability.python.experimental.tangent_spaces.symmetric_matrix import SymmetricMatrixSpace
 
 __all__ = [
     'AxisAlignedSpace',
+    'ConstantDiagonalSymmetricMatrixSpace',
     'FullSpace',
     'GeneralSpace',
     'ProbabilitySimplexSpace',
     'SphericalSpace',
+    'SymmetricMatrixSpace',
     'TangentSpace',
     'UnspecifiedTangentSpaceError',
     'ZeroSpace',

tensorflow_probability/python/experimental/tangent_spaces/simplex.py

@@ -0,0 +1,83 @@
+# Copyright 2023 The TensorFlow Probability Authors.
+#
+# 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.
+# ============================================================================
+
+"""Tangent Spaces related to simplices."""
+
+import numpy as np
+import tensorflow.compat.v2 as tf
+from tensorflow_probability.python.experimental.linalg import linear_operator_row_block as lorb
+from tensorflow_probability.python.experimental.tangent_spaces import spaces
+from tensorflow_probability.python.internal import distribution_util
+from tensorflow_probability.python.internal import dtype_util
+from tensorflow_probability.python.internal import prefer_static as ps
+
+
+class ProbabilitySimplexSpace(spaces.TangentSpace):
+  """Tangent space of M for Simplex distributions in R^n."""
+
+  def compute_basis(self, x):
+    """Returns a `TangentSpace` of a n-simplex."""
+    # The tangent space of the simplex satisfies `{x | <1, x> = 0}`, where `1`
+    # is the vector of all `1`s. This can be seen by the fact that `1` is
+    # orthogonal to the unit simplex.
+    # We can do this by using the basis:  e_i - e_n, 1 <= i <= n - 1. For n = 4,
+    # this looks like:
+    # [[1, 0., 0., -1],
+    #  [0, 1., 0., -1],
+    #  [0, 0., 1., -1]]
+    dim = ps.shape(x)[-1]
+    block1 = tf.linalg.LinearOperatorIdentity(num_rows=dim - 1, dtype=x.dtype)
+    block2 = tf.linalg.LinearOperatorFullMatrix(
+        -tf.ones([dim - 1, 1], dtype=x.dtype))
+    simplex_basis_linop = lorb.LinearOperatorRowBlock([block1, block2])
+    return spaces.LinearOperatorBasis(simplex_basis_linop)
+
+  def _transform_general(self, x, f, **kwargs):
+    basis = self.compute_basis(x)
+    # Note that B @ B.T results in the matrix I + 11^T, where 1 is the vector of
+    # all ones. By the matrix determinant lemma we have det(I + 11^T) = n + 1,
+    # or the dimension of the ambient space.
+    dim = ps.shape(x)[-1]
+    result = dtype_util.as_numpy_dtype(x.dtype)(0.5 * np.log(dim))
+    new_basis_tensor = spaces.compute_new_basis_tensor(f, x, basis)
+    new_log_volume = spaces.volume_coefficient(
+        distribution_util.move_dimension(new_basis_tensor, 0, -2))
+    result = new_log_volume - result
+    return result, spaces.GeneralSpace(
+        spaces.DenseBasis(new_basis_tensor), computed_log_volume=new_log_volume)
+
+  def _transform_coordinatewise(self, x, f, **kwargs):
+    # Compute the diagonal. New matrix is Linop that we can easily write.
+    dim = ps.shape(x)[-1]
+    diag_jacobian = spaces.coordinatewise_jvp(f, x)
+    # Multiplying the basis written in block form as [I, 1] by the diagonal
+    # results in this operator:
+    block1 = tf.linalg.LinearOperatorDiag(diag_jacobian[..., :-1])
+    block2 = tf.linalg.LinearOperatorFullMatrix(
+        diag_jacobian[..., -1:, tf.newaxis] *
+        tf.ones([dim - 1, 1], dtype=x.dtype))
+    linop = lorb.LinearOperatorRowBlock([block1, block2])
+
+    # The volume can be calculated again by the matrix determinant lemma:
+    # det(D**2 + d_n**2 11^T) = (1 + d_n**2 1(D^-1)**21^T) * det(D**2)
+    # = (\sum d_i**-2) * \prod d_i**2
+    log_diag_jacobian = tf.math.log(tf.math.abs(diag_jacobian))
+    log_volume = tf.math.reduce_sum(log_diag_jacobian, axis=-1)
+    log_volume = log_volume + 0.5 * tf.math.reduce_logsumexp(
+        -2. * log_diag_jacobian, axis=-1) - 0.5 * np.log(dim)
+    return log_volume, spaces.GeneralSpace(
+        spaces.LinearOperatorBasis(linop), computed_log_volume=log_volume)
+
+