code: def test_positional_encoder(self, pos_axial_shape):
      # dim should divide FixedBasePositionalEncoding.n_digits
      batch, length, dim = 2, 32, 8
      input_shape = (batch, length, dim)
      vocab_size = 32
      x = np.random.randint(0, vocab_size - 1, input_shape)
      # should sum to dim
      pos_d_axial_embs = (4, 4)
  
      positional_encoding = ct.PositionalEncoder(
          "train",
          dropout=0.1,
          max_len=length,
          pos_axial_shape=pos_axial_shape,
          pos_d_axial_embs=pos_d_axial_embs,
      )
      _, _ = positional_encoding.init(shapes.signature(x))
      y = positional_encoding(x)
      self.assertEqual(y.shape, input_shape)
  
file_path: trax/trax/models/research/configurable_transformer_test.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  from trax import fastmath
  from trax import layers
  from trax import shapes
  from trax.layers import test_utils
  from trax.models.research import configurable_transformer
label_desc: Use image node_id 8 for calling the ConfigurableTransformerTest obj's underlying member method code with example usage: obj.test_positional_encoder(pos_axial_shape) without return types