code: def _GenerateTestData(matrix_shape, num_rhs):
      batch_shape = matrix_shape[:-2]
      matrix_shape = matrix_shape[-2:]
      m = matrix_shape[-2]
      np.random.seed(1)
      matrix = (
          np.random.uniform(
              low=-1.0, high=1.0, size=np.prod(matrix_shape)
          )
          .reshape(matrix_shape)
          .astype(np.float32)
      )
      rhs = np.ones([m, num_rhs]).astype(np.float32)
      matrix = variables.Variable(
          np.tile(matrix, batch_shape + (1, 1)), trainable=False
      )
      rhs = variables.Variable(
          np.tile(rhs, batch_shape + (1, 1)), trainable=False
      )
      return matrix, rhs
  
imports: import numpy
  from tensorflow.python.client import session
  from tensorflow.python.eager import context
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors_impl
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import linalg_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import benchmark
  from tensorflow.python.platform import test
label_desc: Use image node_id 2 for calling a global function with example usage: _GenerateTestData(matrix_shape, num_rhs) and returns: matrix,  rhs