Fix test that failed because of an upstream change.

PiperOrigin-RevId: 584019513
tensorflow · Nov 20, 2023 · 5d09648 · 5d09648
1 parent e0ffe4b
commit 5d09648
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensorflow_probability/python/distributions/linear_gaussian_ssm.py b/tensorflow_probability/python/distributions/linear_gaussian_ssm.py
@@ -61,7 +61,7 @@ def _safe_concat(values):
   for x in values:
     try:
       full_values.append(ps.reshape(x, reference_shape))
-    except (TypeError, ValueError):
+    except (TypeError, ValueError, ZeroDivisionError):
       # JAX/numpy don't like `-1`'s in size-zero shapes.
       full_values.append(ps.reshape(x, trivial_shape))
   return ps.concat(full_values, axis=0)