code: def test_2d_physical(self, dtype, transform_A, transform_B):
      logical_shape = (2, 3, 4)
      A = te.placeholder(shape=logical_shape, dtype=dtype, name="A")
      B = te.compute(
          shape=A.shape, fcompute=lambda i, j, k: A[i, j, k], name="B"
      )
  
      s = te.create_schedule(B.op)
  
      func = self.get_transform(transform_A)
      if func:
          s[A].transform_layout(func)
  
      func = self.get_transform(transform_B)
      if func:
          s[B].transform_layout(func)
  
      # If the two buffers are accessed with the same indices, CSE
      # will replace them with a Let binding.  Since this makes it
      # harder to test what the transformed indices are, disabling
      # the CSE pass for this test.
      with tvm.transform.PassContext(
          disabled_pass=["tir.CommonSubexprElimTIR"]
      ):
          mod = tvm.lower(s, [A, B])
  
      logical_index_vars = self.extract_logical_indices(
          mod["main"].body
      )
      expected_indices_A = self.transform_indices(
          transform_A, logical_shape, logical_index_vars
      )
      expected_indices_B = self.transform_indices(
          transform_B, logical_shape, logical_index_vars
      )
  
      def callback(node):
          if type(node) in [tvm.tir.BufferLoad, tvm.tir.BufferStore]:
              name = node.buffer.name
              if name == "A":
                  expected_indices = expected_indices_A
              elif name == "B":
                  expected_indices = expected_indices_B
              else:
                  raise RuntimeError(f"Unexpected buffer: {name}")
  
              tvm.ir.assert_structural_equal(
                  expected_indices, node.indices
              )
  
      post_order_visit(mod["main"].body, callback)
  
imports: import functools
  import sys
  import pytest
  import numpy
  import tvm
  import tvm.testing
  from tvm import te
  from tvm.tir.stmt_functor import post_order_visit
  from tvm.driver.build_module import schedule_to_module
label_desc: Use image node_id 4 for calling the Test2DPhysicalLayout obj's underlying member method code with example usage: obj.test_2d_physical(dtype, transform_A, transform_B) without return types