Fix bug in pytree flattening of LinearOperatorLowRankUpdate in the JA…

…X backend. PiperOrigin-RevId: 563623848
tensorflow · Sep 8, 2023 · 0fc3b0c · 0fc3b0c
1 parent adbcc1d
commit 0fc3b0c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tensorflow_probability/python/internal/backend/numpy/linalg.py b/tensorflow_probability/python/internal/backend/numpy/linalg.py
@@ -65,7 +65,8 @@ def register_pytrees(env):
           'LinearOperatorScaledIdentity': ('multiplier',),
           'LinearOperatorInversion': ('operator',),
           'LinearOperatorKronecker': ('operators',),
-          'LinearOperatorLowRankUpdate': ('base_operator', 'diag_update'),
+          'LinearOperatorLowRankUpdate': (
+              'base_operator', 'diag_update', 'u', 'v'),
           'LinearOperatorLowerTriangular': ('tril',),
           'LinearOperatorPermutation': ('perm',),
           'LinearOperatorToeplitz': ('col', 'row'),