code: def testDatasetAsNumpy(self):
      arrs = extensions.dataset_as_numpy(
          [tf.constant([1, 2]), tf.constant([3, 4])]
      )
      for a in arrs:
          self.assertIsInstance(a, tf_np.ndarray)
      with self.assertRaisesWithPredicateMatch(
          ValueError,
          r"dataset_as_numpy must be run in eager mode outside tf.function",
      ):
  
          @tf.function
          def f():
              return extensions.dataset_as_numpy([tf.constant([1, 2])])
  
          f()
  
imports: import functools
  import itertools
  from absl import flags
  from absl.testing import parameterized
  from jax import lax
  import numpy
  import tensorflow.compat.v2
  from trax.tf_numpy import extensions
  import trax.tf_numpy.numpy
label_desc: Use image node_id 32 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testDatasetAsNumpy() and returns: extensions