class_docstr: Axial positional encoding.
class_docstr_tok: ["Axial","positional","encoding","."]
code: def init_weights_and_state(self, input_signature):
      d_feature = input_signature.shape[-1]
      if sum(self._d_embs) != d_feature:
          raise ValueError(
              f"sum(self._d_embs) != d_feature: "
              f"sum({self._d_embs}) vs d_feature: {d_feature}"
          )
  
      rngs = fastmath.random.split(self.rng, len(self._d_embs))
      weights = []
      for ax, (ax_rng, d_emb) in enumerate(zip(rngs, self._d_embs)):
          ax_shape = [1] * len(self._shape)
          ax_shape[ax] = self._shape[ax]
          ax_shape = (1,) + tuple(ax_shape) + (d_emb,)
          ax_emb = self._kernel_initializer(ax_shape, ax_rng)
          weights.append(ax_emb)
  
      # State is EMPTY_STATE by default, stays so except for predict mode.
      if self._mode == "predict":
          self.state = np.array(0, dtype=np.int32)
      self.weights = tuple(weights)
  
imports: import logging
  import jax
  import jax.extend
  import numpy
  import trax
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
label_desc: Use image node_id 3 for calling the AxialPositionalEncoding obj's underlying member method code with example usage: obj.init_weights_and_state(input_signature) without return types