code: def testMaxPool3DGradGrad(self, shard_spec, padding):
      np.random.seed(123)
      dep_window_size = 2
      row_window_size = 3
      col_window_size = 4
      window_size = [
          1,
          dep_window_size,
          row_window_size,
          col_window_size,
          1,
      ]
      stride_size = [
          1,
          dep_window_size - 1,
          row_window_size - 1,
          col_window_size - 1,
          1,
      ]
  
      num_deps = 3
      num_rows = (row_window_size - 1) * 5 + 1
      num_cols = (col_window_size - 1) * 7 + 1
      x_in = np.random.normal(
          0.0, 1.0, 8 * num_deps * num_rows * num_cols * 3
      ).reshape([8, num_deps, num_rows, num_cols, 3])
      inputs = constant_op.constant(x_in, dtype=dtypes.float32)
  
      with backprop.GradientTape() as outer_tape:
          with backprop.GradientTape() as inner_tape:
              outer_tape.watch([inputs])
              inner_tape.watch([inputs])
              expected_result = nn_ops.max_pool3d(
                  inputs, window_size, stride_size, padding
              )
          expected_first_grad = inner_tape.gradient(
              expected_result, [inputs]
          )
      expected_second_grad = outer_tape.gradient(
          expected_first_grad, [inputs]
      )
  
      if shard_spec == "sharded":
          layout = self.batch_layout_5d
      else:
          layout = self.replicated_layout_5d
  
      inputs = numpy_util.pack_numpy(inputs, layout)
  
      @polymorphic_function.function()
      def compute_gradients(inputs):
          with backprop.GradientTape() as outer_tape:
              with backprop.GradientTape() as inner_tape:
                  outer_tape.watch([inputs])
                  inner_tape.watch([inputs])
                  dtensor_result = nn_ops.max_pool3d(
                      inputs, window_size, stride_size, padding
                  )
              dtensor_first_grad = inner_tape.gradient(
                  dtensor_result, [inputs]
              )
          dtensor_second_grad = outer_tape.gradient(
              dtensor_first_grad[0], [inputs]
          )
          return dtensor_first_grad, dtensor_second_grad
  
      dtensor_first_grad, dtensor_second_grad = compute_gradients(
          inputs
      )
  
      self.assertDTensorEqual(
          expected_first_grad[0], layout, dtensor_first_grad[0]
      )
      self.assertDTensorEqual(
          expected_second_grad[0], layout, dtensor_second_grad[0]
      )
  
imports: import itertools
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import numpy_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.dtensor.python.tests import test_util_ops
  from tensorflow.python.eager import backprop
  from tensorflow.python.eager import context
  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.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_array_ops
  from tensorflow.python.ops import gen_image_ops
  from tensorflow.python.ops import gen_linalg_ops
  from tensorflow.python.ops import nn_impl
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 15 for calling the DTensorBatchParallelSPMDTest obj's underlying member method code with example usage: obj.testMaxPool3DGradGrad(shard_spec, padding) and returns: dtensor_first_grad,  dtensor_second_grad