code: def testConvWithGradient(
      self, input_shape, kernel_shape, padding, sharding_specs, topology
  ):
      # TODO(b/208700444): add support for SPMD expansion of spatially partitioned
      # conv backprop.
      self.skipTest(
          "b/208700444: Spatially partitioned conv backprop not implemented."
      )
  
      mesh_dims = [spec for spec in sharding_specs if spec != UNSHARDED]
      mesh = self._create_mesh(mesh_dims, topology)
  
      x_in = constant_op.constant(
          np.random.random(size=input_shape), dtype=dtypes.float32
      )
      kernel_in = constant_op.constant(
          np.random.random(size=kernel_shape), dtype=dtypes.float32
      )
  
      @polymorphic_function.function
      def conv_fn(inputs, img_kernel, padding):
          with backprop.GradientTape() as tape:
              tape.watch([inputs, img_kernel])
              output = nn_ops.conv2d_v2(
                  inputs,
                  img_kernel,
                  strides=[1, 1, 1, 1],
                  padding=padding,
              )
          inputs_grad, kernel_grad = tape.gradient(
              output, [inputs, img_kernel]
          )
          return output, inputs_grad, kernel_grad
  
      (
          expected_output,
          expected_inputs_grad,
          expected_kernel_grad,
      ) = conv_fn(x_in, kernel_in, padding)
  
      input_layout = Layout(sharding_specs, mesh)
      kernel_layout = Layout.replicated(mesh, 4)
  
      d_x_in = numpy_util.pack_numpy(x_in, input_layout)
      d_kernel_in = numpy_util.pack_numpy(kernel_in, kernel_layout)
  
      d_output, d_inputs_grad, d_kernel_grad = conv_fn(
          d_x_in, d_kernel_in, padding
      )
  
      self.assertDTensorEqual(expected_output, input_layout, d_output)
      self.assertDTensorEqual(
          expected_inputs_grad, input_layout, d_inputs_grad
      )
      self.assertDTensorEqual(
          expected_kernel_grad, kernel_layout, d_kernel_grad
      )
  
imports: from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import numpy_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.eager import backprop
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import special_math_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 4 for calling the SpatiallyPartitionedConvOpTest obj's underlying member method code with example usage: obj.testConvWithGradient(input_shape, kernel_shape, padding, sharding_specs, topology) and returns: output,  inputs_grad,  kernel_grad