Ignore bad return type annotation in convert_seed_and_sample_shape

Why? This type is incorrect (the type checker does not consider `rng` here to be a `PRNGKey`) The reason this has worked in the past is because `chex.PRNGKey` is considered to be `Any`; this incorrect annotation is blocking us from making `chex.PRNGKey` a more meaningful annotation. PiperOrigin-RevId: 565142299
google-deepmind · Sep 13, 2023 · e0d994d · e0d994d
1 parent 079849a
commit e0d994d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/distrax/_src/distributions/distribution.py b/distrax/_src/distributions/distribution.py
@@ -324,7 +324,7 @@ def convert_seed_and_sample_shape(
   else:  # key is of type PRNGKey
     rng = seed
 
-  return rng, sample_shape
+  return rng, sample_shape  # type: ignore[bad-return-type]
 
 
 def to_batch_shape_index(