code: def fit(A, b, tol, max_iter):
      # Note that this function works even tensors 'A' and 'b' are NumPy or CuPy
      # arrays.
      xp = cupy.get_array_module(A)
      x = xp.zeros_like(b, dtype=np.float64)
      r0 = b - xp.dot(A, x)
      p = r0
      for i in range(max_iter):
          a = xp.inner(r0, r0) / xp.inner(p, xp.dot(A, p))
          x += a * p
          r1 = r0 - a * xp.dot(A, p)
          if xp.linalg.norm(r1) < tol:
              return x
          b = xp.inner(r1, r1) / xp.inner(r0, r0)
          p = r1 + b * p
          r0 = r1
      print("Failed to converge. Increase max-iter or tol.")
      return x
  
imports: import argparse
  import contextlib
  import time
  import numpy
  import cupy
label_desc: Use image node_id 2 for calling a global function with example usage: fit(A, b, tol, max_iter) and returns: x, x