code: def test_can_escape_vectorization_with_none_ndims(self):
      # Suppose the original fn supports `None` as an input.
      fn = lambda x, y: (
          tf.reduce_sum(x, axis=0),
          y[0] if y is not None else y,
      )
  
      polymorphic_fn = vectorization_util.make_rank_polymorphic(
          fn, core_ndims=[1, None]
      )
      rx, ry = polymorphic_fn([[1.0, 2.0, 4.0], [3.0, 5.0, 7.0]], None)
      self.assertAllEqual(rx.shape, [2])
      self.assertIsNone(ry)
  
      single_arg_polymorphic_fn = (
          vectorization_util.make_rank_polymorphic(
              lambda y: fn(tf.convert_to_tensor([1.0, 2.0, 3.0]), y),
              core_ndims=None,
          )
      )
      rx, ry = self.evaluate(
          single_arg_polymorphic_fn(
              tf.convert_to_tensor([[1.0, 3.0], [2.0, 4.0]])
          )
      )
      self.assertAllEqual(ry, [1.0, 3.0])
  
imports: import warnings
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import vectorization_util
label_desc: Use image node_id 8 for calling the VectorizationTest obj's underlying member method code with example usage: obj.test_can_escape_vectorization_with_none_ndims() without return types