parameters: (a,b)
return_statement: return ans
docstring: r"""
      Solve a linear matrix equation Ax = B.This is done by calculating x = A^-1B.So A must not be singular.
      :param a:(...,M,M)
      :param b:(...,M)
      :return:solution of Ax = b formula.x in the shape of (...M)
docstring_summary: r"""
      Solve a linear matrix equation Ax = B.This is done by calculating x = A^-1B.So A must not be singular.
      :param a:(...,M,M)
      :param b:(...,M)
      :return:solution of Ax = b formula.x in the shape of (...M)
docstring_tokens: r Solve a linear matrix equation Ax = B . This is done by calculating x = A^ - 1B . So A must not be singular . : param a : ( ... M M ) : param b : ( ... M ) : return : solution of Ax = b formula . x in the shape of ( ... M )
function: def solve(a,b):
      r"""
      Solve a linear matrix equation Ax = B.This is done by calculating x = A^-1B.So A must not be singular.
      :param a:(...,M,M)
      :param b:(...,M)
      :return:solution of Ax = b formula.x in the shape of (...M)
      """
      def forward_code(np, data):
          a, b = data["inputs"]
          L = data["outputs"][0]
          ans = np.linalg.solve(a, b)
          np.copyto(L, ans)
  
      def backward_code1(np, data):
          def T(x):
              return np.swapaxes(x, -1, -2)
          _dot = partial(np.einsum, '...ij,...jk->...ik')
          dout = data["dout"]
          out = data["outputs"][0]
          f_out = data["f_outputs"][0]
          inp = data["inputs"][0]
          updim = lambda x: x if x.ndim == a.ndim else x[..., None]
          t = -_dot(updim(np.linalg.solve(T(inp), dout)), T(updim(f_out)))
          np.copyto(out, t)
  
      def backward_code2(np, data):
          out = data["outputs"][0]
          np.copyto(out, 0)
  
      l_ans = jt.numpy_code(
          [b.shape],
          [b.dtype],
          [a, b],
          forward_code,
          [backward_code1, backward_code2],
      )
      ans = l_ans[0]
      return ans
function_tokens: def solve ( a , b ) : def forward_code ( np , data ) : a , b = data [ "inputs" ] L = data [ "outputs" ] [ 0 ] ans = np . linalg . solve ( a , b ) np . copyto ( L , ans ) def backward_code1 ( np , data ) : def T ( x ) : return np . swapaxes ( x , - 1 , - 2 ) _dot = partial ( np . einsum , '...ij,...jk->...ik' ) dout = data [ "dout" ] out = data [ "outputs" ] [ 0 ] f_out = data [ "f_outputs" ] [ 0 ] inp = data [ "inputs" ] [ 0 ] updim = lambda x : x if x . ndim == a . ndim else x [ ... , None ] t = - _dot ( updim ( np . linalg . solve ( T ( inp ) , dout ) ) , T ( updim ( f_out ) ) ) np . copyto ( out , t ) def backward_code2 ( np , data ) : out = data [ "outputs" ] [ 0 ] np . copyto ( out , 0 ) l_ans = jt . numpy_code ( [ b . shape ] , [ b . dtype ] , [ a , b ] , forward_code , [ backward_code1 , backward_code2 ] , ) ans = l_ans [ 0 ] return ans