code: def override_shape(tensor_type, axis, dim):
      """Change a dimension in a tensor shape."""
      # Handle multiple tensors by overriding the shape of each.
      if isinstance(tensor_type, relay.TupleType):
          tensor_type = tensor_type.fields
      else:
          tensor_type = [tensor_type]
  
      # Create new tensortypes for each input.
      new_types = []
      for t_type in tensor_type:
          new_dims = list(t_type.shape)
          new_dims[axis] = dim
          new_types.append(relay.TensorType(new_dims, t_type.dtype))
  
      # Dont return a tuple if there is a single tensor.
      if len(new_types) == 1:
          return new_types[0]
      return relay.TupleType(tvm.runtime.convert(new_types))
  
file_path: tvm/python/tvm/relay/transform/flexible_shape.py
imports: import tvm
  from tvm import relay
label_desc: Use image node_id 1 for calling a global function with example usage: override_shape(tensor_type, axis, dim) and returns: relay, new_types