docstr: Test 2d local convolution with circular padding on a 3x3 example
docstr_tok: ["Test","2d","local","convolution","with","circular","padding","on","a","3x3","example"]
code: def test_circular_conv_local_2d_custom(self):
      """
      Test 2d local 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)),
                  ((0, 1, 0), (1, 3, 1), (0, 1, 0)),
                  ((0, 1, 0), (1, 4, 1), (0, 1, 0)),
              ),
              (
                  ((0, 1, 0), (1, 5, 1), (0, 1, 0)),
                  ((0, 1, 0), (1, 6, 1), (0, 1, 0)),
                  ((0, 1, 0), (1, 7, 1), (0, 1, 0)),
              ),
              (
                  ((0, 1, 0), (1, 8, 1), (0, 1, 0)),
                  ((0, 1, 0), (1, 9, 1), (0, 1, 0)),
                  ((0, 1, 0), (1, 10, 1), (0, 1, 0)),
              ),
          )
      )
      kernel = np.expand_dims(kernel, (3,))
      kernel = np.reshape(kernel, (3, 3, 9, 1))
  
      conv_module = nn.ConvLocal(
          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, 9, 1)
      )
      # Compare with manually computed convolution
      correct_ans = np.array(
          (
              (
                  2 * 1 + 7 + 2 + 4 + 3,
                  3 * 2 + 8 + 3 + 5 + 1,
                  4 * 3 + 9 + 1 + 6 + 2,
              ),
              (
                  5 * 4 + 1 + 5 + 7 + 6,
                  6 * 5 + 2 + 6 + 8 + 4,
                  7 * 6 + 3 + 4 + 9 + 5,
              ),
              (
                  8 * 7 + 4 + 8 + 1 + 9,
                  9 * 8 + 5 + 9 + 2 + 7,
                  10 * 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 27 for calling the LinearTest obj's underlying member method code with example usage: obj.test_circular_conv_local_2d_custom() without return types