code: def test_forward(self):
      layer = tl.Fn(
          "SumAndMax",
          lambda x0, x1: (x0 + x1, jnp.maximum(x0, x1)),
          n_out=2,
      )
  
      x0 = np.array([1, 2, 3, 4, 5])
      x1 = np.array([10, 20, 30, 40, 50])
  
      y0, y1 = layer((x0, x1))
      self.assertEqual(y0.tolist(), [11, 22, 33, 44, 55])
      self.assertEqual(y1.tolist(), [10, 20, 30, 40, 50])
  
      y2, y3 = layer.forward((x0, x1))
      self.assertEqual(y2.tolist(), [11, 22, 33, 44, 55])
      self.assertEqual(y3.tolist(), [10, 20, 30, 40, 50])
  
      (y4, y5), state = layer.pure_fn(
          (x0, x1), tl.EMPTY_WEIGHTS, tl.EMPTY_STATE, None
      )
      self.assertEqual(y4.tolist(), [11, 22, 33, 44, 55])
      self.assertEqual(y5.tolist(), [10, 20, 30, 40, 50])
      self.assertEqual(state, tl.EMPTY_STATE)
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  from trax import fastmath
  from trax import shapes
  from trax.fastmath import numpy
  import trax.layers
label_desc: Use image node_id 4 for calling the FnTest obj's underlying member method code with example usage: obj.test_forward() without return types