code: def _all_to_all_grad(op, grad):
      # The gradient of a all-to-all is also a all-to-all but the
      # split_dimension and concat_dimension is swapped.
      # The gradient with respect to group_assignment is None.
      return [
          gen_tpu_ops.all_to_all(
              grad,
              op.inputs[1],
              concat_dimension=op.get_attr("split_dimension"),
              split_dimension=op.get_attr("concat_dimension"),
              split_count=op.get_attr("split_count"),
          ),
          None,
      ]
  
imports: 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_tpu_ops
  from tensorflow.python.ops.gen_tpu_ops import *
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.tpu import tpu_function
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 3 for calling a global function with example usage: _all_to_all_grad(op, grad) and returns: list