docstr: Securely computes xa * wa + xb * wb, where:
      1. xa is a tensor that belongs to rank_a, and xb is a tensor that belongs to rank_b.
      2. wa is a shared tensor for rank_a and rank_b.
      3. wb is a shared tensor for rank_a and rank_b.
  
      Input:    xa   wa.share_a   wb.share_a       xb   wa.share_b   wb.share_b
  
      Comp:       (xa * wa.share_b).share_a          (xa * wa.share_b).share_b         (SMM)
                  (xb * wa.share_a).share_a          (xb * wa.share_a).share_b         (SMM)
                      xa * wa.share_a                       xb * wb.share_b
  
      Output:  z.share_a = (xa * wa.share_b).share_a + (xb * wb.share_a).share_a + (xa * wb.share_a)
               z.share_b = (xa * wa.share_b).share_b + (xb * wb.share_a).share_b + (xb * wa.share_b)
docstr_tok: ["Securely","computes","xa","*","wa","+","xb","*","wb",",","where",":","1.","xa","is","a","tensor","that","belongs","to","rank_a",",","and","xb","is","a","tensor","that","belongs","to","rank_b",".","2.","wa","is","a","shared","tensor","for","rank_a","and","rank_b",".","3.","wb","is","a","shared","tensor","for","rank_a","and","rank_b",".","Input",":","xa","wa.share_a","wb.share_a","xb","wa.share_b","wb.share_b","Comp",":","(","xa","*","wa.share_b",")",".share_a","(","xa","*","wa.share_b",")",".share_b","(","SMM",")","(","xb","*","wa.share_a",")",".share_a","(","xb","*","wa.share_a",")",".share_b","(","SMM",")","xa","*","wa.share_a","xb","*","wb.share_b","Output",":","z.share_a","=","(","xa","*","wa.share_b",")",".share_a","+","(","xb","*","wb.share_a",")",".share_a","+","(","xa","*","wb.share_a",")","z.share_b","=","(","xa","*","wa.share_b",")",".share_b","+","(","xb","*","wb.share_a",")",".share_b","+","(","xb","*","wa.share_b",")"]
code: def cross_smm(
      cls,
      ctx: Context,
      group,
      xa,
      xb,
      wa: ArithmeticSharedTensor,
      wb: ArithmeticSharedTensor,
      rank_a,
      rank_b,
      phe_cipher,
      precision_bits=None,
  ):
      """
      Securely computes xa * wa + xb * wb, where:
          1. xa is a tensor that belongs to rank_a, and xb is a tensor that belongs to rank_b.
          2. wa is a shared tensor for rank_a and rank_b.
          3. wb is a shared tensor for rank_a and rank_b.
  
          Input:    xa   wa.share_a   wb.share_a       xb   wa.share_b   wb.share_b
  
          Comp:       (xa * wa.share_b).share_a          (xa * wa.share_b).share_b         (SMM)
                      (xb * wa.share_a).share_a          (xb * wa.share_a).share_b         (SMM)
                          xa * wa.share_a                       xb * wb.share_b
  
          Output:  z.share_a = (xa * wa.share_b).share_a + (xb * wb.share_a).share_a + (xa * wb.share_a)
                   z.share_b = (xa * wa.share_b).share_b + (xb * wb.share_a).share_b + (xb * wa.share_b)
      """
      from fate.arch.protocol.mpc.mpc import FixedPointEncoder
      from fate.arch.context import PHECipher
  
      assert isinstance(wa, ArithmeticSharedTensor), "invalid wa"
      assert isinstance(wb, ArithmeticSharedTensor), "invalid wb"
      assert (
          isinstance(rank_a, int) and 0 <= rank_a < ctx.world_size
      ), f"invalid rank_a: {rank_a}"
      assert (
          isinstance(rank_b, int) and 0 <= rank_b < ctx.world_size
      ), f"invalid rank_b: {rank_b}"
      assert isinstance(phe_cipher, PHECipher), "invalid phe_cipher"
      encoder = FixedPointEncoder(precision_bits)
  
      x = ctx.mpc.cond_call(lambda: xa, lambda: xb, dst=rank_a)
      encoded_x = encoder.encode(x)
      w = ctx.mpc.cond_call(lambda: wa, lambda: wb, dst=rank_a)
      with IgnoreEncodings([w]):
          z = w.rmatmul(encoded_x)
  
      # [xa|rank_a] @ [wa.share|rank_b]
      z += cls.smm(
          ctx,
          group,
          op=lambda a, b: torch.matmul(b, a),
          rank_1=rank_b,
          tensor_1=ctx.mpc.option(wa.share, rank_b),
          cipher_1=ctx.mpc.option(phe_cipher, rank_b),
          rank_2=rank_a,
          tensor_2=ctx.mpc.option(encoded_x, rank_a),
      )
      # [xb|rank_b] @ [wb.share|rank_a]
      z += cls.smm(
          ctx,
          group,
          op=lambda a, b: torch.matmul(b, a),
          rank_1=rank_a,
          tensor_1=ctx.mpc.option(wb.share, rank_a),
          cipher_1=ctx.mpc.option(phe_cipher, rank_a),
          rank_2=rank_b,
          tensor_2=ctx.mpc.option(encoded_x, rank_b),
      )
      with IgnoreEncodings([z]):
          z.div_(encoder.scale)
      return z
  
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 1 for calling the SSHE obj's underlying member method code with example usage: obj.cross_smm(cls, ctx, group, xa, xb, wa, wb, rank_a, rank_b, phe_cipher, precision_bits) and returns: z