code: def test_expectation_penalty(profiler):
      penalty = ExpectationProfilerPenalty(profiler, 20000, 2.0)
      loss, details = penalty(42.0, {"conv": {0: 0.2, 1: 0.3, 2: 0.5}})
      assert details["loss_original"] == 42.0
      assert details["penalty"] == 33800
      assert details["normalized_penalty"] == 0.69
      assert loss == 42.0 + 0.69 * 2.0
  
      prob = torch.tensor([0.2, 0.3, 0.5])
      loss, details = penalty(
          torch.tensor(42.0), {"conv": {i: prob[i] for i in range(3)}}
      )
      assert abs(loss.item() - (42.0 + 0.69 * 2.0)) < 1e-4
  
      penalty = ExpectationProfilerPenalty(
          profiler, 40000, -2.0, nonlinear="negative"
      )
      loss, details = penalty(42.0, {"conv": {0: 0.2, 1: 0.3, 2: 0.5}})
      assert details["normalized_penalty"] == 33800 / 40000 - 1
      assert loss == 42.0 + 0.155 * 2.0
  
      loss, details = penalty(42.0, {"conv": {0: 0.0, 1: 0.0, 2: 1.0}})
      assert details["normalized_penalty"] == 0
      assert loss == 42.0
  
      penalty = ExpectationProfilerPenalty(
          profiler, 40000, 2.0, nonlinear="absolute", aggregate="mul"
      )
      loss, details = penalty(42.0, {"conv": {0: 0.2, 1: 0.3, 2: 0.5}})
      assert details["normalized_penalty"] == abs(33800 / 40000 - 1)
      assert loss == 42.0 * (1 + 0.155) ** 2
  
      penalty = ExpectationProfilerPenalty(
          profiler, 30000, 2.0, nonlinear="positive", aggregate="mul"
      )
      loss, details = penalty(42.0, {"conv": {0: 0.2, 1: 0.3, 2: 0.5}})
      assert details["normalized_penalty"] == 33800 / 30000 - 1
      assert loss == 42.0 * (33800 / 30000) ** 2
  
      loss, details = penalty(42.0, {"conv": {0: 1.0, 1: 0.0, 2: 0.0}})
      assert details["normalized_penalty"] == 0
      assert loss == 42.0
  
file_path: nni/test/algo/nas/oneshot/test_profiler.py
imports: import pytest
  import torch
  from torch.nn import Conv2d
  from nni.nas.nn import ModelSpace, LayerChoice
  from nni.nas.oneshot.pytorch.profiler import RangeProfilerFilter, ExpectationProfilerPenalty, SampleProfilerPenalty
  from nni.nas.profiler.pytorch.flops import FlopsProfiler
label_desc: Use image node_id 3 for calling a global function with example usage: test_expectation_penalty(profiler) without return types