code: def test_all_to_all_split_count_not_equal_to_group_assignment_shape(
      self,
  ):
      with self.assertRaisesRegex(
          ValueError,
          "split_count 1 must equal the size of the second dimension "
          "of group_assignment 2",
      ):
          tpu_ops.all_to_all(
              x=[0.0, 0.1652, 0.6543],
              group_assignment=[[0, 1], [2, 3]],
              concat_dimension=0,
              split_dimension=0,
              split_count=1,
          )
  
imports: from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import importer
  from tensorflow.python.framework import ops
  from tensorflow.python.layers import convolutional
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import control_flow_util
  from tensorflow.python.ops import init_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import special_math_ops
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.platform import test
  from tensorflow.python.tpu import tpu
  from tensorflow.python.tpu import tpu_feed
  from tensorflow.python.tpu import tpu_replication
  from tensorflow.python.tpu import training_loop
  from tensorflow.python.tpu.ops import tpu_ops
label_desc: Use image node_id 3 for calling the TPUOpsTest obj's underlying member method code with example usage: obj.test_all_to_all_split_count_not_equal_to_group_assignment_shape() without return types