code: def rename_dimension(x, old_name, new_name):
    
    return reshape(x, x.shape.rename_dimension(old_name, new_name))
docs: Reshape a Tensor, renaming one dimension.
  
    Args:
      x: a Tensor
      old_name: a string
      new_name: a string
  
    Returns:
      a Tensor