code: def test_create_device_mesh_for_nd_torus(
      self, devices, mesh_shape, expected_assignment
  ):
      jax_devices = devices(True)
      physical_mesh = mesh_utils._get_physical_tpu_mesh(jax_devices)
      _, assignment = mesh_utils._create_device_mesh_for_nd_torus(
          physical_mesh, mesh_shape
      )
      self.assertEqual(assignment, expected_assignment)
  
imports: import collections
  from collections.abc import Sequence
  import dataclasses
  import numpy
  from absl import logging
  from absl.testing import absltest
  from absl.testing import parameterized
  from jax.experimental import mesh_utils
  from jax.sharding import Mesh
  from jax._src import test_util
label_desc: Use image node_id 3 for calling the MeshUtilsTest obj's underlying member method code with example usage: obj.test_create_device_mesh_for_nd_torus(devices, mesh_shape, expected_assignment) without return types