code: def square(x):
      (x_shares,) = replicate_shares([x.share])
      x1, x2 = x_shares
  
      x_square = x1**2 + 2 * x1 * x2
  
      z = x.shallow_copy()
      z.share = x_square
      return z
  
file_path: FATE/python/fate/arch/protocol/mpc/primitives/replicated.py
imports: import torch
  from ..None import communicator
label_desc: Use image node_id 9 for calling a global function with example usage: square(x) and returns: z