code: def testSlicewiseOperationAndGenericGradOperation(self):
      slicewise_operation = mtf.SlicewiseOperation(
          tf.exp,
          [self.x],
          [self.x.shape],
          [self.x.dtype],
          splittable_dims=[
              self.a_dim
          ],  # pretend only dim "a" can be split.
          grad_function=lambda op, dy: [dy * op.outputs[0]],
          name="component-wise exp",
      )
  
      self.assertEqual(
          slicewise_operation.splittable_dims, frozenset(["a"])
      )
      self.assertEqual(
          slicewise_operation.unsplittable_dims, frozenset(["b"])
      )
  
      generic_grad_operation = mtf.GenericGradOperation(
          slicewise_operation, [self.x]
      )
  
      self.assertEqual(
          generic_grad_operation.splittable_dims, frozenset(["a", "b"])
      )
      self.assertEqual(
          generic_grad_operation.unsplittable_dims, frozenset()
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import mesh_tensorflow
  from mesh_tensorflow import test_utils
  import numpy
  import tensorflow.compat.v1
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 3 for calling the OperationSplittabilityTest obj's underlying member method code with example usage: obj.testSlicewiseOperationAndGenericGradOperation() without return types