code: def test_batched_lbfgs_rosenbrock():
      """Test batched rosenbrock using batched lbfgs implemtnation"""
  
      num_batches = 5
      np.random.seed(42)
      a = np.random.normal(1, scale=0.1, size=num_batches)
      b = np.random.normal(100, scale=10, size=num_batches)
  
      def f(x, n=None):
          nonlocal a
          nonlocal b
          nonlocal num_batches
  
          if n is not None:
              return rosenbrock(x, a[n], b[n])
  
          fb = batched_rosenbrock(x, num_batches, a, b)
          return fb
  
      def gf(x, n=None):
          nonlocal a
          nonlocal b
          nonlocal num_batches
  
          if n is not None:
              return g_rosenbrock(x, a[n], b[n])
  
          g = g_batched_rosenbrock(x, num_batches, a, b)
          return g
  
      x0 = np.zeros(2 * num_batches)
      x0[0] = 0.0
      x0[1] = 0.0
  
      # analytical minimum
      res_true = np.zeros(num_batches * 2)
      for i in range(num_batches):
          res_true[i * 2 : (i + 1) * 2] = np.array([a[i], a[i] ** 2])
  
      # our new batch-aware l-bfgs optimizer
      res_xk, _, _ = batched_fmin_lbfgs_b(
          f, x0, num_batches, gf, iprint=-1, factr=100
      )
  
      np.testing.assert_allclose(res_xk, res_true, rtol=1e-5)
  
imports: from cuml.tsa.batched_lbfgs import batched_fmin_lbfgs_b
  from cuml.internals.safe_imports import cpu_only_import
label_desc: Use image node_id 5 for calling a global function with example usage: test_batched_lbfgs_rosenbrock() and returns: fb, g, rosenbrock, g_rosenbrock