old_contents: import unittest
  
  import numpy as np
  
  import chainer
  from chainer import backend
  import chainer.initializers as I
  from chainer import optimizer_hooks
  from chainer import optimizers
  from chainer import testing
  from chainer.testing import attr
  
  
  class SimpleLink(chainer.Link):
  
      def __init__(self, w, g):
          super(SimpleLink, self).__init__()
          with self.init_scope():
              self.param = chainer.Parameter(I.Zero(), w.shape)
          self.param.data = w
          self.param.grad = g
  
  
  class TestLasso(unittest.TestCase):
      def setUp(self):
          self.target = SimpleLink(
              np.arange(6, dtype=np.float32).reshape(2, 3),
              np.arange(3, -3, -1, dtype=np.float32).reshape(2, 3))
  
      def check_lasso(self):
          w = self.target.param.data
          g = self.target.param.grad
          xp = backend.get_array_module(w)
          decay = 0.2
          expect = w - g - decay * xp.sign(w)
  
          opt = optimizers.SGD(lr=1)
          opt.setup(self.target)
          opt.add_hook(optimizer_hooks.Lasso(decay))
          opt.update()
  
          testing.assert_allclose(expect, w)
  
      def test_lasso_cpu(self):
          self.check_lasso()
  
      @attr.gpu
      def test_lasso_gpu(self):
          self.target.to_gpu()
          self.check_lasso()
  
  
  testing.run_module(__name__, __file__)
  
new_contents: import unittest
  
  import numpy as np
  
  import chainer
  from chainer import optimizer_hooks
  from chainer import optimizers
  from chainer import testing
  
  
  _backend_params = [
      # NumPy
      {},
      # {'use_ideep': 'always'},
      # CuPy
      {'use_cuda': True, 'cuda_device': 0},
      # {'use_cuda': True, 'cuda_device': 1},
      # ChainerX
      {'use_chainerx': True, 'chainerx_device': 'native:0'},
      {'use_chainerx': True, 'chainerx_device': 'cuda:0'},
      # {'use_chainerx': True, 'chainerx_device': 'cuda:1'},
  ]
  
  
  class SimpleLink(chainer.Link):
  
      def __init__(self, params):
          super(SimpleLink, self).__init__()
          with self.init_scope():
              for i, p in enumerate(params):
                  setattr(self, 'p{}'.format(i), p)
  
  
  @testing.backend.inject_backend_tests(None, _backend_params)
  @testing.backend.inject_backend_tests(None, _backend_params)
  @testing.backend.inject_backend_tests(None, _backend_params)
  class TestLasso(unittest.TestCase):
      def setUp(self):
          num_params = 3
          arrs = [
              np.random.uniform(-3, 3, (2, 3)).astype(np.float32)
              for _ in range(num_params)]
          grads = [
              np.random.uniform(-3, 3, (2, 3)).astype(np.float32)
              for _ in range(num_params)]
          params = []
          for arr, grad in zip(arrs, grads):
              param = chainer.Parameter(arr)
              param.grad = grad
              params.append(param)
  
          self.target = SimpleLink(params)
  
      def check_lasso(self, backend_configs):
          target = self.target
          assert len(backend_configs) == len(list(target.params()))
          devices = [bc.device for bc in backend_configs]
  
          decay = 0.2
  
          expects = []
          # Compute expected
          for param, device in zip(target.params(), devices):
              expects.append(param.array - param.grad -
                             decay * np.sign(param.array))
              param.to_device(device)
  
          # Compute using optimizer_hook
          opt = optimizers.SGD(lr=1)
          opt.setup(self.target)
          opt.add_hook(optimizer_hooks.Lasso(decay))
          opt.update()
  
          # Validate
          for expect, param in zip(expects, target.params()):
              testing.assert_allclose(expect, param.array)
  
      def test_lasso_cpu(self, backend_config0,
                         backend_config1, backend_config2):
          self.check_lasso([backend_config0, backend_config1, backend_config2])
  
      @testing.attr.gpu
      def test_lasso_gpu(self, backend_config0,
                         backend_config1, backend_config2):
          self.target.to_gpu()
          self.check_lasso([backend_config0, backend_config1, backend_config2])
  
  
  testing.run_module(__name__, __file__)
  
subject: update test to also check with chainerx