docstr: Securely computes x * x, where:
      1. x is a shared tensor that belongs to rank_a and rank_b.
      2. cipher_a is a PHE cipher that belongs to rank_a.
docstr_tok: ["Securely","computes","x","*","x",",","where",":","1.","x","is","a","shared","tensor","that","belongs","to","rank_a","and","rank_b",".","2.","cipher_a","is","a","PHE","cipher","that","belongs","to","rank_a","."]
code: def mpc_square(
      cls,
      ctx: Context,
      group,
      rank_a,
      rank_b,
      x: ArithmeticSharedTensor,
      cipher_a,
  ):
      """
      Securely computes x * x, where:
          1. x is a shared tensor that belongs to rank_a and rank_b.
          2. cipher_a is a PHE cipher that belongs to rank_a.
      """
      x_a_enc = ctx.mpc.communicator.broadcast_obj(
          obj=ctx.mpc.option_call(
              lambda: cipher_a.get_tensor_encryptor().encrypt_tensor(
                  x.share
              ),
              dst=rank_a,
          ),
          src=rank_a,
      )
      x_ab = cls.phe_to_mpc(
          ctx,
          group,
          rank_a=rank_b,
          rank_b=rank_a,
          phe_cipher_b=ctx.mpc.option(cipher_a, dst=rank_a),
          phe_tensor_a=ctx.mpc.option_call(
              lambda: x_a_enc * x.share, dst=rank_b
          ),
      )
      x_ab.share = x.share * x.share + x_ab.share * 2
      with IgnoreEncodings([x_ab]):
          x_ab.div_(x.encoder.scale)
      return x_ab
  
file_path: FATE/python/fate/arch/protocol/mpc/primitives/sshe.py
imports: import typing
  import torch
  from fate.arch.context import Context
  from fate.arch.tensor import DTensor
  from fate.arch.protocol.mpc.common.rng import generate_random_ring_element
  from fate.arch.protocol.mpc.primitives.arithmetic import ArithmeticSharedTensor
  from fate.arch.protocol.mpc.primitives.beaver import IgnoreEncodings
  from fate.arch.trace import auto_trace
label_desc: Use image node_id 6 for calling the SSHE obj's underlying member method code with example usage: obj.mpc_square(cls, ctx, group, rank_a, rank_b, x, cipher_a) and returns: x_ab