code: def assertDTypesEqual(self, a, b):
      get_dtype = lambda t: t.dtype
      self.assertEqual(
          tf.nest.map_structure(get_dtype, a),
          tf.nest.map_structure(get_dtype, b),
      )
  
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 21 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.assertDTypesEqual(a, b) without return types