jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def ort_validate(self, onnx_io, inputs, outputs, tolerate_small_mismatch=False):

        inputs, _ = torch.jit._flatten(inputs)
        outputs, _ = torch.jit._flatten(outputs)

        def to_numpy(tensor):
            if tensor.requires_grad:
                return tensor.detach().cpu().numpy()
            else:
                return tensor.cpu().numpy()

        inputs = list(map(to_numpy, inputs))
        outputs = list(map(to_numpy, outputs))

        ort_session = onnxruntime.InferenceSession(onnx_io.getvalue())
        # compute onnxruntime output prediction
        ort_inputs = dict((ort_session.get_inputs()[i].name, inpt) for i, inpt in enumerate(inputs))
        ort_outs = ort_session.run(None, ort_inputs)
        for i, element in enumerate(outputs):
            try:
                torch.testing.assert_close(element, ort_outs[i], rtol=1e-03, atol=1e-05)
            except AssertionError as error:
                if tolerate_small_mismatch:
                    self.assertIn("(0.00%)", str(error), str(error))
                else:
                    raise
//  [python]
def test_inference_mask_generation_three_boxes_point_batch(self):
        model = SamModel.from_pretrained("facebook/sam-vit-base")
        processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

        model.to(torch_device)
        model.eval()

        raw_image = prepare_image()

        # fmt: off
        input_boxes = torch.Tensor([[[620, 900, 1000, 1255]], [[75, 275, 1725, 850]],  [[75, 275, 1725, 850]]]).cpu()
        EXPECTED_IOU = torch.tensor([[[0.9773, 0.9881, 0.9522],
         [0.5996, 0.7661, 0.7937],
         [0.5996, 0.7661, 0.7937]]])
        # fmt: on
        input_boxes = input_boxes.unsqueeze(0)

        inputs = processor(raw_image, input_boxes=input_boxes, return_tensors="pt").to(torch_device)

        with torch.no_grad():
            outputs = model(**inputs)

        iou_scores = outputs.iou_scores.cpu()
        self.assertTrue(iou_scores.shape == (1, 3, 3))
        torch.testing.assert_allclose(iou_scores, EXPECTED_IOU, atol=1e-4, rtol=1e-4)
//  [python]
def test_inference_mask_generation_three_boxes_point_batch(self):
        model = SamModel.from_pretrained("facebook/sam-vit-base")
        processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

        model.to(torch_device)
        model.eval()

        raw_image = prepare_image()

        # fmt: off
        input_boxes = torch.Tensor([[[620, 900, 1000, 1255]], [[75, 275, 1725, 850]],  [[75, 275, 1725, 850]]]).cpu()
        EXPECTED_IOU = torch.tensor([[[0.9773, 0.9881, 0.9522],
         [0.5996, 0.7661, 0.7937],
         [0.5996, 0.7661, 0.7937]]])
        # fmt: on
        input_boxes = input_boxes.unsqueeze(0)

        inputs = processor(raw_image, input_boxes=input_boxes, return_tensors="pt").to(torch_device)

        with torch.no_grad():
            outputs = model(**inputs)

        iou_scores = outputs.iou_scores.cpu()
        self.assertTrue(iou_scores.shape == (1, 3, 3))
        torch.testing.assert_close(iou_scores, EXPECTED_IOU, atol=1e-4, rtol=1e-4)
//  [python]
def test_compute_loss(self) -> None:
        """
        Test loss computation
        """
        prediction = Trajectory(data=self.prediction_data)
        target = Trajectory(data=self.target_data)
        scenarios = [CachedScenario(log_name='', token='lane_following_with_lead', scenario_type='') for _ in range(2)]

        loss = self.objective.compute(
            {"trajectory": prediction.to_feature_tensor()}, {"trajectory": target.to_feature_tensor()}, scenarios
        )
        torch.testing.assert_allclose(loss, torch.tensor(0.60653, dtype=torch.float64))
//  [python]
def test_compute_loss(self) -> None:
        """
        Test loss computation
        """
        prediction = Trajectory(data=self.prediction_data)
        target = Trajectory(data=self.target_data)
        scenarios = [CachedScenario(log_name='', token='lane_following_with_lead', scenario_type='') for _ in range(2)]

        loss = self.objective.compute(
            {"trajectory": prediction.to_feature_tensor()}, {"trajectory": target.to_feature_tensor()}, scenarios
        )
        torch.testing.assert_close(loss, torch.tensor(0.60653, dtype=torch.float64))
//  [python]
def test_to_rgb(self):
        cfg = dict(type='ClsDataPreprocessor', to_rgb=True)
        processor: ClsDataPreprocessor = MODELS.build(cfg)

        data = {'inputs': [torch.randint(0, 256, (3, 224, 224))]}
        inputs = processor(data)['inputs']
        torch.testing.assert_allclose(data['inputs'][0].flip(0).float(),
                                      inputs[0])

        data = {'inputs': torch.randint(0, 256, (1, 3, 224, 224))}
        inputs = processor(data)['inputs']
        torch.testing.assert_allclose(data['inputs'].flip(1).float(), inputs)
//  [python]
def test_to_rgb(self):
        cfg = dict(type='ClsDataPreprocessor', to_rgb=True)
        processor: ClsDataPreprocessor = MODELS.build(cfg)

        data = {'inputs': [torch.randint(0, 256, (3, 224, 224))]}
        inputs = processor(data)['inputs']
        torch.testing.assert_close(data['inputs'][0].flip(0).float(),
                                      inputs[0])

        data = {'inputs': torch.randint(0, 256, (1, 3, 224, 224))}
        inputs = processor(data)['inputs']
        torch.testing.assert_close(data['inputs'].flip(1).float(), inputs)
//  [python]
def test_call(self, trajectory):
        transformer = MeanFunction(DeltaState())
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            torch.testing.assert_allclose(transformed_observation.state, obs.state)
            torch.testing.assert_allclose(transformed_observation.action, obs.action)
            torch.testing.assert_allclose(transformed_observation.reward, obs.reward)
            assert transformed_observation.done == obs.done

            mean = observation.state
            torch.testing.assert_allclose(
                transformed_observation.next_state, obs.next_state - mean
            )
//  [python]
def test_call(self, trajectory):
        transformer = MeanFunction(DeltaState())
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            torch.testing.assert_close(transformed_observation.state, obs.state)
            torch.testing.assert_close(transformed_observation.action, obs.action)
            torch.testing.assert_close(transformed_observation.reward, obs.reward)
            assert transformed_observation.done == obs.done

            mean = observation.state
            torch.testing.assert_close(
                transformed_observation.next_state, obs.next_state - mean
            )
//  [python]
def test_custom_mean(self, trajectory):
        transformer = MeanFunction(MeanModel())
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            torch.testing.assert_allclose(transformed_observation.state, obs.state)
            torch.testing.assert_allclose(transformed_observation.action, obs.action)
            torch.testing.assert_allclose(transformed_observation.reward, obs.reward)
            assert transformed_observation.done == obs.done

            mean = 2 * obs.state + obs.action
            torch.testing.assert_allclose(
                transformed_observation.next_state, obs.next_state - mean
            )
//  [python]
def test_custom_mean(self, trajectory):
        transformer = MeanFunction(MeanModel())
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            torch.testing.assert_close(transformed_observation.state, obs.state)
            torch.testing.assert_close(transformed_observation.action, obs.action)
            torch.testing.assert_close(transformed_observation.reward, obs.reward)
            assert transformed_observation.done == obs.done

            mean = 2 * obs.state + obs.action
            torch.testing.assert_close(
                transformed_observation.next_state, obs.next_state - mean
            )
//  [python]
def test_layer_norm_op_pytorch(self, X, eps, elementwise_affine, gc, dc):
        axis = np.random.randint(0, len(X.shape))

        if elementwise_affine:
            gamma = np.random.randn(*X.shape[axis:]).astype(np.float32)
            beta = np.random.randn(*X.shape[axis:]).astype(np.float32)
            expected_norm, expected_mean, expected_std = \
                _layer_norm_with_affine_ref(axis, eps, X, gamma, beta)
            actual_norm, actual_mean, actual_std = torch.ops._caffe2.LayerNorm(
                torch.tensor(X), torch.tensor(gamma), torch.tensor(beta),
                axis, eps, True)
        else:
            expected_norm, expected_mean, expected_std = _layer_norm_ref(
                axis, eps, X)
            actual_norm, actual_mean, actual_std = torch.ops._caffe2.LayerNorm(
                torch.tensor(X), None, None, axis, eps)

        torch.testing.assert_allclose(
            expected_norm, actual_norm, rtol=1e-4, atol=1e-4)
        torch.testing.assert_allclose(expected_mean, actual_mean)
        torch.testing.assert_allclose(expected_std, actual_std)
//  [python]
def test_layer_norm_op_pytorch(self, X, eps, elementwise_affine, gc, dc):
        axis = np.random.randint(0, len(X.shape))

        if elementwise_affine:
            gamma = np.random.randn(*X.shape[axis:]).astype(np.float32)
            beta = np.random.randn(*X.shape[axis:]).astype(np.float32)
            expected_norm, expected_mean, expected_std = \
                _layer_norm_with_affine_ref(axis, eps, X, gamma, beta)
            actual_norm, actual_mean, actual_std = torch.ops._caffe2.LayerNorm(
                torch.tensor(X), torch.tensor(gamma), torch.tensor(beta),
                axis, eps, True)
        else:
            expected_norm, expected_mean, expected_std = _layer_norm_ref(
                axis, eps, X)
            actual_norm, actual_mean, actual_std = torch.ops._caffe2.LayerNorm(
                torch.tensor(X), None, None, axis, eps)

        torch.testing.assert_close(
            expected_norm, actual_norm, rtol=1e-4, atol=1e-4)
        torch.testing.assert_close(expected_mean, actual_mean)
        torch.testing.assert_close(expected_std, actual_std)
//  [python]
def test_layer_norm_op_pytorch_cuda(self, X, eps, elementwise_affine):
        axis = np.random.randint(0, len(X.shape))

        if elementwise_affine:
            gamma = np.random.randn(*X.shape[axis:]).astype(np.float32)
            beta = np.random.randn(*X.shape[axis:]).astype(np.float32)
            expected_norm, expected_mean, expected_std = \
                _layer_norm_with_affine_ref(axis, eps, X, gamma, beta)
            actual_norm, actual_mean, actual_std = torch.ops._caffe2.LayerNorm(
                torch.tensor(X).cuda(),
                torch.tensor(gamma).cuda(),
                torch.tensor(beta).cuda(),
                axis,
                eps,
                True)
        else:
            expected_norm, expected_mean, expected_std = _layer_norm_ref(
                axis, eps, X)
            actual_norm, actual_mean, actual_std = torch.ops._caffe2.LayerNorm(
                torch.tensor(X).cuda(), None, None, axis, eps)

        torch.testing.assert_allclose(
            expected_norm, actual_norm.cpu(), rtol=1e-4, atol=1e-4)
        torch.testing.assert_allclose(expected_mean, actual_mean.cpu())
        torch.testing.assert_allclose(expected_std, actual_std.cpu())
//  [python]
def test_layer_norm_op_pytorch_cuda(self, X, eps, elementwise_affine):
        axis = np.random.randint(0, len(X.shape))

        if elementwise_affine:
            gamma = np.random.randn(*X.shape[axis:]).astype(np.float32)
            beta = np.random.randn(*X.shape[axis:]).astype(np.float32)
            expected_norm, expected_mean, expected_std = \
                _layer_norm_with_affine_ref(axis, eps, X, gamma, beta)
            actual_norm, actual_mean, actual_std = torch.ops._caffe2.LayerNorm(
                torch.tensor(X).cuda(),
                torch.tensor(gamma).cuda(),
                torch.tensor(beta).cuda(),
                axis,
                eps,
                True)
        else:
            expected_norm, expected_mean, expected_std = _layer_norm_ref(
                axis, eps, X)
            actual_norm, actual_mean, actual_std = torch.ops._caffe2.LayerNorm(
                torch.tensor(X).cuda(), None, None, axis, eps)

        torch.testing.assert_close(
            expected_norm, actual_norm.cpu(), rtol=1e-4, atol=1e-4)
        torch.testing.assert_close(expected_mean, actual_mean.cpu())
        torch.testing.assert_close(expected_std, actual_std.cpu())
//  [python]
def test_layer_norm_op_jit(self, X, eps, elementwise_affine, gc, dc):
        @torch.jit.script
        def jit_layer_norm(X, gamma=None, beta=None, axis=1, eps=1e-5,
                           elementwise_affine=False):
            # type: (Tensor, Optional[Tensor], Optional[Tensor], int, float, bool) -> Tuple[Tensor, Tensor, Tensor]
            return torch.ops._caffe2.LayerNorm(
                X, gamma, beta, axis, eps, elementwise_affine)

        axis = np.random.randint(0, len(X.shape))

        if elementwise_affine:
            gamma = np.random.randn(*X.shape[axis:]).astype(np.float32)
            beta = np.random.randn(*X.shape[axis:]).astype(np.float32)
            expected_norm, expected_mean, expected_std = \
                _layer_norm_with_affine_ref(axis, eps, X, gamma, beta)
            actual_norm, actual_mean, actual_std = jit_layer_norm(
                torch.Tensor(X), torch.tensor(gamma), torch.tensor(beta),
                axis, eps, elementwise_affine)
        else:
            expected_norm, expected_mean, expected_std = _layer_norm_ref(
                axis, eps, X)
            actual_norm, actual_mean, actual_std = jit_layer_norm(
                torch.tensor(X), None, None, axis, eps, elementwise_affine)

        torch.testing.assert_allclose(
            expected_norm, actual_norm, rtol=1e-4, atol=1e-4)
        torch.testing.assert_allclose(expected_mean, actual_mean)
        torch.testing.assert_allclose(expected_std, actual_std)
//  [python]
def test_layer_norm_op_jit(self, X, eps, elementwise_affine, gc, dc):
        @torch.jit.script
        def jit_layer_norm(X, gamma=None, beta=None, axis=1, eps=1e-5,
                           elementwise_affine=False):
            # type: (Tensor, Optional[Tensor], Optional[Tensor], int, float, bool) -> Tuple[Tensor, Tensor, Tensor]
            return torch.ops._caffe2.LayerNorm(
                X, gamma, beta, axis, eps, elementwise_affine)

        axis = np.random.randint(0, len(X.shape))

        if elementwise_affine:
            gamma = np.random.randn(*X.shape[axis:]).astype(np.float32)
            beta = np.random.randn(*X.shape[axis:]).astype(np.float32)
            expected_norm, expected_mean, expected_std = \
                _layer_norm_with_affine_ref(axis, eps, X, gamma, beta)
            actual_norm, actual_mean, actual_std = jit_layer_norm(
                torch.Tensor(X), torch.tensor(gamma), torch.tensor(beta),
                axis, eps, elementwise_affine)
        else:
            expected_norm, expected_mean, expected_std = _layer_norm_ref(
                axis, eps, X)
            actual_norm, actual_mean, actual_std = jit_layer_norm(
                torch.tensor(X), None, None, axis, eps, elementwise_affine)

        torch.testing.assert_close(
            expected_norm, actual_norm, rtol=1e-4, atol=1e-4)
        torch.testing.assert_close(expected_mean, actual_mean)
        torch.testing.assert_close(expected_std, actual_std)
//  [python]
def test_inverse(self, trajectory):
        transformer = MeanFunction(DeltaState())
        for observation in trajectory:
            obs = observation.clone()
            inverse_observation = transformer.inverse(transformer(observation))
            for x, y in zip(obs, inverse_observation):
                torch.testing.assert_allclose(x, y)
//  [python]
def test_inverse(self, trajectory):
        transformer = MeanFunction(DeltaState())
        for observation in trajectory:
            obs = observation.clone()
            inverse_observation = transformer.inverse(transformer(observation))
            for x, y in zip(obs, inverse_observation):
                torch.testing.assert_close(x, y)
//  [python]
def test_layer_norm_op_c10_preallocated_outputs(
            self, X, eps, elementwise_affine, gc, dc):
        # This test case ensures that it works correctly when output tensors are
        # preallocated.
        axis = np.random.randint(0, len(X.shape))

        self.ws.create_blob("X").feed(X)
        if elementwise_affine:
            gamma = np.random.randn(*X.shape[axis:]).astype(np.float32)
            beta = np.random.randn(*X.shape[axis:]).astype(np.float32)
            self.ws.create_blob("gamma").feed(gamma)
            self.ws.create_blob("beta").feed(beta)

        m = ModelHelper(name="test")
        m.net.C10LayerNorm_DontUseThisOpYet(
            ["X", "gamma", "beta"] if elementwise_affine else ["X"],
            ["Y", "mean", "std"],
            axis=axis,
            epsilon=eps,
            elementwise_affine=elementwise_affine,
        )
        self.ws.create_net(m.param_init_net).run()
        net = self.ws.create_net(m.net)
        # run two times to be extra sure that the outputs are preallocated
        net.run()
        net.run()

        if elementwise_affine:
            expected_norm, expected_mean, expected_std = \
                _layer_norm_with_affine_ref(axis, eps, X, gamma, beta)
        else:
            expected_norm, expected_mean, expected_std = _layer_norm_ref(
                axis, eps, X)
        actual_norm = self.ws.fetch_blob('Y')
        actual_mean = self.ws.fetch_blob('mean')
        actual_std = self.ws.fetch_blob('std')

        torch.testing.assert_allclose(
            expected_norm, actual_norm, rtol=1e-4, atol=1e-4)
        torch.testing.assert_allclose(expected_mean, actual_mean)
        torch.testing.assert_allclose(expected_std, actual_std)
//  [python]
def test_layer_norm_op_c10_preallocated_outputs(
            self, X, eps, elementwise_affine, gc, dc):
        # This test case ensures that it works correctly when output tensors are
        # preallocated.
        axis = np.random.randint(0, len(X.shape))

        self.ws.create_blob("X").feed(X)
        if elementwise_affine:
            gamma = np.random.randn(*X.shape[axis:]).astype(np.float32)
            beta = np.random.randn(*X.shape[axis:]).astype(np.float32)
            self.ws.create_blob("gamma").feed(gamma)
            self.ws.create_blob("beta").feed(beta)

        m = ModelHelper(name="test")
        m.net.C10LayerNorm_DontUseThisOpYet(
            ["X", "gamma", "beta"] if elementwise_affine else ["X"],
            ["Y", "mean", "std"],
            axis=axis,
            epsilon=eps,
            elementwise_affine=elementwise_affine,
        )
        self.ws.create_net(m.param_init_net).run()
        net = self.ws.create_net(m.net)
        # run two times to be extra sure that the outputs are preallocated
        net.run()
        net.run()

        if elementwise_affine:
            expected_norm, expected_mean, expected_std = \
                _layer_norm_with_affine_ref(axis, eps, X, gamma, beta)
        else:
            expected_norm, expected_mean, expected_std = _layer_norm_ref(
                axis, eps, X)
        actual_norm = self.ws.fetch_blob('Y')
        actual_mean = self.ws.fetch_blob('mean')
        actual_std = self.ws.fetch_blob('std')

        torch.testing.assert_close(
            expected_norm, actual_norm, rtol=1e-4, atol=1e-4)
        torch.testing.assert_close(expected_mean, actual_mean)
        torch.testing.assert_close(expected_std, actual_std)
//  [python]
def test_call(self, trajectory):
        transformer = RewardScaler(scale=5.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            torch.testing.assert_allclose(transformed_observation.state, obs.state)
            torch.testing.assert_allclose(transformed_observation.action, obs.action)
            torch.testing.assert_allclose(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
            torch.testing.assert_allclose(
                transformed_observation.reward, obs.reward / 5
            )
//  [python]
def test_call(self, trajectory):
        transformer = RewardScaler(scale=5.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            torch.testing.assert_close(transformed_observation.state, obs.state)
            torch.testing.assert_close(transformed_observation.action, obs.action)
            torch.testing.assert_close(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
            torch.testing.assert_close(
                transformed_observation.reward, obs.reward / 5
            )
//  [python]
def test_inverse(self, trajectory):
        transformer = RewardClipper(min_reward=0.0, max_reward=1.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            inverse_observation = transformer.inverse(transformed_observation)

            torch.testing.assert_allclose(transformed_observation.state, obs.state)
            torch.testing.assert_allclose(transformed_observation.action, obs.action)
            torch.testing.assert_allclose(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
            torch.testing.assert_allclose(
                inverse_observation.reward, transformed_observation.reward
            )
//  [python]
def test_inverse(self, trajectory):
        transformer = RewardClipper(min_reward=0.0, max_reward=1.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            inverse_observation = transformer.inverse(transformed_observation)

            torch.testing.assert_close(transformed_observation.state, obs.state)
            torch.testing.assert_close(transformed_observation.action, obs.action)
            torch.testing.assert_close(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
            torch.testing.assert_close(
                inverse_observation.reward, transformed_observation.reward
            )
//  [python]
def test_inverse(self, trajectory):
        transformer = ActionClipper(min_action=-1000.0, max_action=1000.0)

        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            inverse_observation = transformer.inverse(transformed_observation)

            torch.testing.assert_allclose(transformed_observation.state, obs.state)
            torch.testing.assert_allclose(transformed_observation.action, obs.action)
            torch.testing.assert_allclose(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
            torch.testing.assert_allclose(
                inverse_observation.reward, transformed_observation.reward
            )
//  [python]
def test_inverse(self, trajectory):
        transformer = ActionClipper(min_action=-1000.0, max_action=1000.0)

        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            inverse_observation = transformer.inverse(transformed_observation)

            torch.testing.assert_close(transformed_observation.state, obs.state)
            torch.testing.assert_close(transformed_observation.action, obs.action)
            torch.testing.assert_close(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
            torch.testing.assert_close(
                inverse_observation.reward, transformed_observation.reward
            )
//  [python]
def test_call(self, trajectory):
        transformer = ActionClipper(min_action=0.0, max_action=1.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            assert (transformed_observation.action >= 0.0).all()
            assert (transformed_observation.action <= 1.0).all()
            torch.testing.assert_allclose(transformed_observation.state, obs.state)
            torch.testing.assert_allclose(transformed_observation.reward, obs.reward)
            torch.testing.assert_allclose(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
//  [python]
def test_call(self, trajectory):
        transformer = ActionClipper(min_action=0.0, max_action=1.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            assert (transformed_observation.action >= 0.0).all()
            assert (transformed_observation.action <= 1.0).all()
            torch.testing.assert_close(transformed_observation.state, obs.state)
            torch.testing.assert_close(transformed_observation.reward, obs.reward)
            torch.testing.assert_close(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
//  [python]
def test_call(self, trajectory):
        transformer = RewardClipper(min_reward=0.0, max_reward=1.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            assert transformed_observation.reward >= 0.0
            assert transformed_observation.reward <= 1.0
            torch.testing.assert_allclose(transformed_observation.state, obs.state)
            torch.testing.assert_allclose(transformed_observation.action, obs.action)
            torch.testing.assert_allclose(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done

            if 0 <= obs.reward <= 1:
                assert transformed_observation.reward == obs.reward
            elif obs.reward <= 0:
                assert transformed_observation.reward == 0.0
            elif obs.reward >= 1.0:
                assert transformed_observation.reward == 1.0
//  [python]
def test_call(self, trajectory):
        transformer = RewardClipper(min_reward=0.0, max_reward=1.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            assert transformed_observation.reward >= 0.0
            assert transformed_observation.reward <= 1.0
            torch.testing.assert_close(transformed_observation.state, obs.state)
            torch.testing.assert_close(transformed_observation.action, obs.action)
            torch.testing.assert_close(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done

            if 0 <= obs.reward <= 1:
                assert transformed_observation.reward == obs.reward
            elif obs.reward <= 0:
                assert transformed_observation.reward == 0.0
            elif obs.reward >= 1.0:
                assert transformed_observation.reward == 1.0
//  [python]
def test_inverse(self, trajectory):
        transformer = RewardScaler(scale=5.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            inverse_observation = transformer.inverse(transformed_observation)

            torch.testing.assert_allclose(transformed_observation.state, obs.state)
            torch.testing.assert_allclose(transformed_observation.action, obs.action)
            torch.testing.assert_allclose(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
            torch.testing.assert_allclose(inverse_observation.reward, obs.reward)
//  [python]
def test_inverse(self, trajectory):
        transformer = RewardScaler(scale=5.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            inverse_observation = transformer.inverse(transformed_observation)

            torch.testing.assert_close(transformed_observation.state, obs.state)
            torch.testing.assert_close(transformed_observation.action, obs.action)
            torch.testing.assert_close(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
            torch.testing.assert_close(inverse_observation.reward, obs.reward)
//  [python]
def test_call(self, trajectory):
        transformer = ActionScaler(scale=5.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            torch.testing.assert_allclose(transformed_observation.state, obs.state)
            torch.testing.assert_allclose(transformed_observation.reward, obs.reward)
            torch.testing.assert_allclose(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
            torch.testing.assert_allclose(
                transformed_observation.action, obs.action / 5
            )
//  [python]
def test_call(self, trajectory):
        transformer = ActionScaler(scale=5.0)
        for observation in trajectory:
            obs = observation.clone()
            transformed_observation = transformer(observation)
            torch.testing.assert_close(transformed_observation.state, obs.state)
            torch.testing.assert_close(transformed_observation.reward, obs.reward)
            torch.testing.assert_close(
                transformed_observation.next_state, obs.next_state
            )
            assert transformed_observation.done == obs.done
            torch.testing.assert_close(
                transformed_observation.action, obs.action / 5
            )
//  [python]
def test_update(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = ActionNormalizer(
            dim=(trajectory.action.shape[-1],), preserve_origin=preserve_origin
        )

        mean = torch.mean(trajectory.action, 0)
        var = torch.var(trajectory.action, 0)

        transformer.update(trajectory)
        torch.testing.assert_allclose(transformer._normalizer.mean, mean)
        torch.testing.assert_allclose(transformer._normalizer.variance, var)
//  [python]
def test_update(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = ActionNormalizer(
            dim=(trajectory.action.shape[-1],), preserve_origin=preserve_origin
        )

        mean = torch.mean(trajectory.action, 0)
        var = torch.var(trajectory.action, 0)

        transformer.update(trajectory)
        torch.testing.assert_close(transformer._normalizer.mean, mean)
        torch.testing.assert_close(transformer._normalizer.variance, var)
//  [python]
def test_inverse(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = ActionNormalizer(
            dim=(trajectory.action.shape[-1],), preserve_origin=preserve_origin
        )
        transformer.update(trajectory)

        observation = get_observation()
        obs = observation.clone()
        inverse_observation = transformer.inverse(transformer(observation))
        for x, y in zip(obs, inverse_observation):
            if x.shape == y.shape:
                torch.testing.assert_allclose(x, y)
//  [python]
def test_inverse(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = ActionNormalizer(
            dim=(trajectory.action.shape[-1],), preserve_origin=preserve_origin
        )
        transformer.update(trajectory)

        observation = get_observation()
        obs = observation.clone()
        inverse_observation = transformer.inverse(transformer(observation))
        for x, y in zip(obs, inverse_observation):
            if x.shape == y.shape:
                torch.testing.assert_close(x, y)
//  [python]
def test_update(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = StateNormalizer(
            dim=(trajectory.state.shape[-1],), preserve_origin=preserve_origin
        )

        mean = torch.mean(trajectory.state, 0)
        var = torch.var(trajectory.state, 0)

        transformer.update(trajectory)
        torch.testing.assert_allclose(transformer._normalizer.mean, mean)
        torch.testing.assert_allclose(transformer._normalizer.variance, var)
//  [python]
def test_update(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = StateNormalizer(
            dim=(trajectory.state.shape[-1],), preserve_origin=preserve_origin
        )

        mean = torch.mean(trajectory.state, 0)
        var = torch.var(trajectory.state, 0)

        transformer.update(trajectory)
        torch.testing.assert_close(transformer._normalizer.mean, mean)
        torch.testing.assert_close(transformer._normalizer.variance, var)
//  [python]
def test_inverse(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = StateNormalizer(
            dim=(trajectory.state.shape[-1],), preserve_origin=preserve_origin
        )

        transformer.update(trajectory)

        observation = get_observation()
        obs = observation.clone()
        inverse_observation = transformer.inverse(transformer(observation))
        for x, y in zip(obs, inverse_observation):
            if x.shape == y.shape:
                torch.testing.assert_allclose(x, y)
//  [python]
def test_inverse(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = StateNormalizer(
            dim=(trajectory.state.shape[-1],), preserve_origin=preserve_origin
        )

        transformer.update(trajectory)

        observation = get_observation()
        obs = observation.clone()
        inverse_observation = transformer.inverse(transformer(observation))
        for x, y in zip(obs, inverse_observation):
            if x.shape == y.shape:
                torch.testing.assert_close(x, y)
//  [python]
def assertTensorEqual(self,
                          tensor: torch.Tensor,
                          value: float,
                          msg=None,
                          **kwarg):
        tensor = tensor.to(torch.float32)
        if tensor.dim() == 0:
            tensor = tensor.unsqueeze(0)
        value = torch.FloatTensor([value])
        try:
            torch.testing.assert_allclose(tensor, value, **kwarg)
        except AssertionError as e:
            self.fail(self._formatMessage(msg, str(e) + str(tensor)))
//  [python]
def assertTensorEqual(self,
                          tensor: torch.Tensor,
                          value: float,
                          msg=None,
                          **kwarg):
        tensor = tensor.to(torch.float32)
        if tensor.dim() == 0:
            tensor = tensor.unsqueeze(0)
        value = torch.FloatTensor([value])
        try:
            torch.testing.assert_close(tensor, value, **kwarg)
        except AssertionError as e:
            self.fail(self._formatMessage(msg, str(e) + str(tensor)))
//  [python]
def test_call(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = StateNormalizer(
            dim=(trajectory.state.shape[-1],), preserve_origin=preserve_origin
        )

        transformer.update(trajectory)
        observation = get_observation()
        obs = observation.clone()
        transformed = transformer(observation)

        if preserve_origin:
            mean = 0
            scale = torch.sqrt(
                transformer._normalizer.variance + transformer._normalizer.mean ** 2
            )
        else:
            mean = transformer._normalizer.mean
            scale = torch.sqrt(transformer._normalizer.variance)
        torch.testing.assert_allclose(transformed.state, (obs.state - mean) / scale)
        torch.testing.assert_allclose(transformed.next_state, obs.next_state)

        torch.testing.assert_allclose(transformed.action, obs.action)
        torch.testing.assert_allclose(transformed.reward, obs.reward)
        assert transformed.done == obs.done
//  [python]
def test_call(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = StateNormalizer(
            dim=(trajectory.state.shape[-1],), preserve_origin=preserve_origin
        )

        transformer.update(trajectory)
        observation = get_observation()
        obs = observation.clone()
        transformed = transformer(observation)

        if preserve_origin:
            mean = 0
            scale = torch.sqrt(
                transformer._normalizer.variance + transformer._normalizer.mean ** 2
            )
        else:
            mean = transformer._normalizer.mean
            scale = torch.sqrt(transformer._normalizer.variance)
        torch.testing.assert_close(transformed.state, (obs.state - mean) / scale)
        torch.testing.assert_close(transformed.next_state, obs.next_state)

        torch.testing.assert_close(transformed.action, obs.action)
        torch.testing.assert_close(transformed.reward, obs.reward)
        assert transformed.done == obs.done
//  [python]
def test_call(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = ActionNormalizer(
            dim=(trajectory.action.shape[-1],), preserve_origin=preserve_origin
        )

        transformer.update(trajectory)
        observation = get_observation()
        obs = observation.clone()
        transformed = transformer(observation)

        if preserve_origin:
            mean = 0
            scale = torch.sqrt(
                transformer._normalizer.variance + transformer._normalizer.mean ** 2
            )
        else:
            mean = transformer._normalizer.mean
            scale = torch.sqrt(transformer._normalizer.variance)
        torch.testing.assert_allclose(transformed.action, (obs.action - mean) / scale)

        torch.testing.assert_allclose(transformed.state, obs.state)
        torch.testing.assert_allclose(transformed.reward, obs.reward)
        torch.testing.assert_allclose(transformed.next_state, obs.next_state)
        assert transformed.done == obs.done
//  [python]
def test_call(self, trajectory, preserve_origin):
        trajectory = stack_list_of_tuples(trajectory)

        transformer = ActionNormalizer(
            dim=(trajectory.action.shape[-1],), preserve_origin=preserve_origin
        )

        transformer.update(trajectory)
        observation = get_observation()
        obs = observation.clone()
        transformed = transformer(observation)

        if preserve_origin:
            mean = 0
            scale = torch.sqrt(
                transformer._normalizer.variance + transformer._normalizer.mean ** 2
            )
        else:
            mean = transformer._normalizer.mean
            scale = torch.sqrt(transformer._normalizer.variance)
        torch.testing.assert_close(transformed.action, (obs.action - mean) / scale)

        torch.testing.assert_close(transformed.state, obs.state)
        torch.testing.assert_close(transformed.reward, obs.reward)
        torch.testing.assert_close(transformed.next_state, obs.next_state)
        assert transformed.done == obs.done
//  [python]
def test_atomref(model_name, enable_atomref):
    dataset = DummyDataset(has_atomref=True)
    atomref = Atomref(max_z=100, dataset=dataset, enable=enable_atomref)
    z, pos, batch = create_example_batch()

    # create model with atomref
    pl.seed_everything(1234)
    model_atomref = create_model(
        load_example_args(model_name, prior_model="Atomref"), prior_model=atomref
    )
    # create model without atomref
    pl.seed_everything(1234)
    model_no_atomref = create_model(load_example_args(model_name, remove_prior=True))

    # get output from both models
    x_atomref, _ = model_atomref(z, pos, batch)
    x_no_atomref, _ = model_no_atomref(z, pos, batch)

    # check if the output of both models differs by the expected atomref contribution
    if enable_atomref:
        expected_offset = scatter(dataset.get_atomref().squeeze()[z], batch).unsqueeze(1)
    else:
        expected_offset = 0
    torch.testing.assert_allclose(x_atomref, x_no_atomref + expected_offset)
//  [python]
def test_atomref(model_name, enable_atomref):
    dataset = DummyDataset(has_atomref=True)
    atomref = Atomref(max_z=100, dataset=dataset, enable=enable_atomref)
    z, pos, batch = create_example_batch()

    # create model with atomref
    pl.seed_everything(1234)
    model_atomref = create_model(
        load_example_args(model_name, prior_model="Atomref"), prior_model=atomref
    )
    # create model without atomref
    pl.seed_everything(1234)
    model_no_atomref = create_model(load_example_args(model_name, remove_prior=True))

    # get output from both models
    x_atomref, _ = model_atomref(z, pos, batch)
    x_no_atomref, _ = model_no_atomref(z, pos, batch)

    # check if the output of both models differs by the expected atomref contribution
    if enable_atomref:
        expected_offset = scatter(dataset.get_atomref().squeeze()[z], batch).unsqueeze(1)
    else:
        expected_offset = 0
    torch.testing.assert_close(x_atomref, x_no_atomref + expected_offset)
//  [python]
def test_batch(self):
        x = torch.tensor([[[0.3, 0.6], [0.5, 0.9]], [[0.6, 0.3], [0.6, 0.9]]]).log()
        decoder = CTCGreedyDecoder()
        r = decoder(x, segmentation=True, apply_softmax=False)
        e = [[1], [1]]
        self.assertEqual(e, r["hyp"])
        # note: checking with ctc_loss does not work for every x
        loss = torch.nn.functional.ctc_loss(
            x,
            torch.tensor(e),
            torch.tensor([1, 1]),
            torch.tensor([1, 1]),
            reduction="none",
        )
        e = loss.neg().exp()
        r = torch.tensor([p.mean() for p in r["prob-segmentation"]])
        torch.testing.assert_allclose(r, e)
//  [python]
def test_batch(self):
        x = torch.tensor([[[0.3, 0.6], [0.5, 0.9]], [[0.6, 0.3], [0.6, 0.9]]]).log()
        decoder = CTCGreedyDecoder()
        r = decoder(x, segmentation=True, apply_softmax=False)
        e = [[1], [1]]
        self.assertEqual(e, r["hyp"])
        # note: checking with ctc_loss does not work for every x
        loss = torch.nn.functional.ctc_loss(
            x,
            torch.tensor(e),
            torch.tensor([1, 1]),
            torch.tensor([1, 1]),
            reduction="none",
        )
        e = loss.neg().exp()
        r = torch.tensor([p.mean() for p in r["prob-segmentation"]])
        torch.testing.assert_close(r, e)
//  [python]
def test_scatter_topk_1d(device):
    src = torch.tensor([
        2.3, 8.3, 1.4,
        3.7, 6.6, 1.7, 8.6, 6.4, 7.6,
        6.8, 9.1,
        5.2, 2.0, 0.4, 1.0
    ], device=device)
    index = torch.tensor([
        0, 0, 0,
        1, 1, 1, 1, 1, 1,
        2, 2,
        3, 3, 3, 3
    ], device=device)

    result_values, result_indexes_whole, result_indexes_within = scatter_topk_1d(src, index, k=4)

    expected_values = torch.tensor([
        [8.3, 2.3, 1.4, float('NaN')],
        [8.6, 7.6, 6.6, 6.4],
        [9.1, 6.8, float('NaN'), float('NaN')],
        [5.2, 2.0, 1.0, 0.4],
    ], device=device)
    expected_indexes_whole = torch.tensor([
        [ 1,  0,  2, -1],
        [ 6,  8,  4,  7],
        [10,  9, -1, -1],
        [11, 12, 14, 13],
    ], device=device)
    expected_indexes_within = torch.tensor([
        [1,  0,  2, -1],
        [3,  5,  1,  4],
        [1,  0, -1, -1],
        [0,  1,  3,  2],
    ], device=device)

    torch.testing.assert_allclose(result_values, expected_values, equal_nan=True)
    torch.testing.assert_allclose(result_indexes_whole, expected_indexes_whole)
    torch.testing.assert_allclose(result_indexes_within, expected_indexes_within)
//  [python]
def test_scatter_topk_1d(device):
    src = torch.tensor([
        2.3, 8.3, 1.4,
        3.7, 6.6, 1.7, 8.6, 6.4, 7.6,
        6.8, 9.1,
        5.2, 2.0, 0.4, 1.0
    ], device=device)
    index = torch.tensor([
        0, 0, 0,
        1, 1, 1, 1, 1, 1,
        2, 2,
        3, 3, 3, 3
    ], device=device)

    result_values, result_indexes_whole, result_indexes_within = scatter_topk_1d(src, index, k=4)

    expected_values = torch.tensor([
        [8.3, 2.3, 1.4, float('NaN')],
        [8.6, 7.6, 6.6, 6.4],
        [9.1, 6.8, float('NaN'), float('NaN')],
        [5.2, 2.0, 1.0, 0.4],
    ], device=device)
    expected_indexes_whole = torch.tensor([
        [ 1,  0,  2, -1],
        [ 6,  8,  4,  7],
        [10,  9, -1, -1],
        [11, 12, 14, 13],
    ], device=device)
    expected_indexes_within = torch.tensor([
        [1,  0,  2, -1],
        [3,  5,  1,  4],
        [1,  0, -1, -1],
        [0,  1,  3,  2],
    ], device=device)

    torch.testing.assert_close(result_values, expected_values, equal_nan=True)
    torch.testing.assert_close(result_indexes_whole, expected_indexes_whole)
    torch.testing.assert_close(result_indexes_within, expected_indexes_within)
//  [python]
def test():
    # the format of list([k1, v1, k2, v2, ...]), ensure config[::2] is key and config[1::2] is value
    model_config = [
        ("BACKEND", "tensorrt"),
        ("SCHEDULE.FUSE_BN", True),
        ("W.QSCHEME", "per-channel-symmetric"),
        ("W.QUANTIZER.TYPE", "uniform"),
        ("W.QUANTIZER.BIT", 8),
        ("W.OBSERVER.TYPE", "MINMAX"),
        ("A.QSCHEME", "per-tensor-symmetric"),
        ("A.QUANTIZER.TYPE", "uniform"),
        ("A.QUANTIZER.BIT", 8),
        ("A.OBSERVER.TYPE", "MINMAX"),
        ("A.OBSERVER.LAYOUT", "NCHW"),
    ]
    model_config = [j for i in model_config for j in i]

    model = torch.hub.load(
        "facebookresearch/deit:main", "deit_tiny_patch16_224", pretrained=True
    )
    config = build_config(model_config)
    qmodel = QuantModel(model, config)

    data = torch.randn(2, 3, 224, 224)
    model.eval()
    qmodel.eval()
    out1 = model(data)
    out2 = qmodel(data)
    torch.testing.assert_allclose(out1, out2, atol=1e-4, rtol=1e-4)
//  [python]
def test():
    # the format of list([k1, v1, k2, v2, ...]), ensure config[::2] is key and config[1::2] is value
    model_config = [
        ("BACKEND", "tensorrt"),
        ("SCHEDULE.FUSE_BN", True),
        ("W.QSCHEME", "per-channel-symmetric"),
        ("W.QUANTIZER.TYPE", "uniform"),
        ("W.QUANTIZER.BIT", 8),
        ("W.OBSERVER.TYPE", "MINMAX"),
        ("A.QSCHEME", "per-tensor-symmetric"),
        ("A.QUANTIZER.TYPE", "uniform"),
        ("A.QUANTIZER.BIT", 8),
        ("A.OBSERVER.TYPE", "MINMAX"),
        ("A.OBSERVER.LAYOUT", "NCHW"),
    ]
    model_config = [j for i in model_config for j in i]

    model = torch.hub.load(
        "facebookresearch/deit:main", "deit_tiny_patch16_224", pretrained=True
    )
    config = build_config(model_config)
    qmodel = QuantModel(model, config)

    data = torch.randn(2, 3, 224, 224)
    model.eval()
    qmodel.eval()
    out1 = model(data)
    out2 = qmodel(data)
    torch.testing.assert_close(out1, out2, atol=1e-4, rtol=1e-4)