code: def test_composite_learning_rate_op(
      self, gc, min_num_iter, max_num_iter
  ):
      np.random.seed(65535)
      # Generate the iteration numbers for sub policy
      # The four sub policies are as follows:
      # 1. exp; 2. step; 3. fix; 4. exp
      num_lr_policy = 4
      iter_nums = np.random.randint(
          low=min_num_iter, high=max_num_iter, size=num_lr_policy
      )
      accu_iter_num = copy.deepcopy(iter_nums)
      for i in range(1, num_lr_policy):
          accu_iter_num[i] += accu_iter_num[i - 1]
      total_iter_nums = accu_iter_num[-1]
  
      policy_lr_scale = np.random.uniform(
          low=0.1, high=2.0, size=num_lr_policy
      )
  
      # args for StepLRPolicy
      step_size = np.random.randint(low=2, high=min_num_iter // 2)
      step_gamma = np.random.random()
      # args for ExpLRPolicy
      exp_gamma = np.random.random()
      # common args
      base_lr = 0.1
  
      # StepLRPolicy
      def step_lr(iter, lr_scale):
          return math.pow(step_gamma, iter // step_size) * lr_scale
  
      # ExpLRPolicy
      def exp_lr(iter, lr_scale):
          return math.pow(exp_gamma, iter) * lr_scale
  
      # FixedLRPolicy
      def fixed_lr(iter, lr_scale):
          return lr_scale
  
      # test one sub policy case
      def one_policy_check_ref(iter, lr_scale):
          iter = int(iter)
          exp_lr_val = exp_lr(iter, lr_scale=lr_scale)
          return (np.array(base_lr * exp_lr_val),)
  
      op = core.CreateOperator(
          "LearningRate",
          "data",
          "out",
          policy="composite",
          sub_policy_num_iters=iter_nums[:1],
          sub_policy_0_lr_scale=policy_lr_scale[0],
          sub_policy_0_policy="exp",
          sub_policy_0_gamma=exp_gamma,
          base_lr=base_lr,
      )
      for iter_idx in range(1, total_iter_nums + 1):
          self.assertReferenceChecks(
              gc,
              op,
              [np.asarray([iter_idx])],
              partial(
                  one_policy_check_ref, lr_scale=policy_lr_scale[0]
              ),
          )
  
      # all the case with all four sub policies
      def all_sub_policy_check_ref(iter, lr_scale):
          assert iter <= accu_iter_num[3]
          if iter <= accu_iter_num[0]:
              lr = exp_lr(iter, lr_scale=lr_scale)
          elif iter <= accu_iter_num[1]:
              lr = step_lr(iter, lr_scale=lr_scale)
          elif iter <= accu_iter_num[2]:
              lr = fixed_lr(iter, lr_scale=lr_scale)
          else:
              lr = exp_lr(iter, lr_scale=lr_scale)
          return (np.array(base_lr * lr),)
  
      op = core.CreateOperator(
          "LearningRate",
          "data",
          "out",
          policy="composite",
          sub_policy_num_iters=iter_nums,
          sub_policy_0_policy="exp",
          sub_policy_0_lr_scale=policy_lr_scale[0],
          sub_policy_0_gamma=exp_gamma,
          sub_policy_1_policy="step",
          sub_policy_1_lr_scale=policy_lr_scale[1],
          sub_policy_1_stepsize=step_size,
          sub_policy_1_gamma=step_gamma,
          sub_policy_2_policy="fixed",
          sub_policy_2_lr_scale=policy_lr_scale[2],
          sub_policy_3_policy="exp",
          sub_policy_3_gamma=exp_gamma,
          sub_policy_3_lr_scale=policy_lr_scale[3],
          base_lr=base_lr,
      )
  
      iter_policy = 0
      for iter_idx in range(1, total_iter_nums + 1):
          if iter_idx > accu_iter_num[iter_policy]:
              iter_policy += 1
          self.assertReferenceChecks(
              gc,
              op,
              [np.asarray([iter_idx])],
              partial(
                  all_sub_policy_check_ref,
                  lr_scale=policy_lr_scale[iter_policy],
              ),
          )
  
file_path: pytorch/caffe2/python/operator_test/learning_rate_op_test.py
imports: from caffe2.python import core
  import caffe2.python.hypothesis_test_util
  import caffe2.python.serialized_test.serialized_test_util
  from hypothesis import given, settings
  import hypothesis.strategies
  import copy
  from functools import partial
  import math
  import numpy
label_desc: Use image node_id 5 for calling the TestLearningRate obj's underlying member method code with example usage: obj.test_composite_learning_rate_op(gc, min_num_iter, max_num_iter) and returns: unknown, unknown, lr_scale, np, np