code: def get_channel_dim(shape, data_format=INVALID):
      assert data_format != INVALID
      assert len(shape) == 4
      if data_format == "NHWC":
          return int(shape[3])
      elif data_format == "NCHW":
          return int(shape[1])
      else:
          raise ValueError("Not a valid data_format", data_format)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling a global function with example usage: get_channel_dim(shape, data_format) and returns: int, int