docstr: Test 2d convolution with circular padding on a 3x3 example.
docstr_tok: ["Test","2d","convolution","with","circular","padding","on","a","3x3","example","."]
code: def test_circular_conv_2d_custom(self):
      """Test 2d convolution with circular padding on a 3x3 example."""
      rng = dict(params=random.key(0))
      x = np.array(((1, 2, 3), (4, 5, 6), (7, 8, 9)))
      x = np.expand_dims(x, (0, 3))
      kernel = np.array(((0, 1, 0), (1, 2, 1), (0, 1, 0)))
      kernel = np.expand_dims(kernel, (2, 3))
  
      conv_module = nn.Conv(
          features=1,
          kernel_size=(3, 3),
          padding="CIRCULAR",
          kernel_init=lambda *_: kernel,
          bias_init=initializers.zeros,
      )
      y, initial_params = conv_module.init_with_output(rng, x)
  
      self.assertEqual(
          initial_params["params"]["kernel"].shape, (3, 3, 1, 1)
      )
      # Compare with manually computed convolution
      correct_ans = np.array(
          (
              (
                  2 * 1 + 7 + 2 + 4 + 3,
                  2 * 2 + 8 + 3 + 5 + 1,
                  2 * 3 + 9 + 1 + 6 + 2,
              ),
              (
                  2 * 4 + 1 + 5 + 7 + 6,
                  2 * 5 + 2 + 6 + 8 + 4,
                  2 * 6 + 3 + 4 + 9 + 5,
              ),
              (
                  2 * 7 + 4 + 8 + 1 + 9,
                  2 * 8 + 5 + 9 + 2 + 7,
                  2 * 9 + 6 + 7 + 3 + 8,
              ),
          )
      )
      correct_ans = np.expand_dims(correct_ans, (0, 3))
      np.testing.assert_allclose(y, correct_ans)
  
imports: import functools
  import jax
  import jax.numpy
  import numpy
  from absl.testing import absltest, parameterized
  from jax import random
  from jax.nn import initializers
  from flax import linen
label_desc: Use image node_id 26 for calling the LinearTest obj's underlying member method code with example usage: obj.test_circular_conv_2d_custom() without return types