code: def square(input: DTensor):
      return DTensor(input.shardings.map_shard(torch.square))
  
imports: import torch
  from ._tensor import DTensor, implements
label_desc: Use image node_id 3 for calling a global function with example usage: square(input) and returns: DTensor