class_docstr: Axial positional encoding.
class_docstr_tok: ["Axial","positional","encoding","."]
code: def forward(self, inputs):
      rng, state = self.rng, self.state
      embs = []
      for ax_emb in self.weights:
          ax_emb = jnp.broadcast_to(
              ax_emb,
              (inputs.shape[0],) + self._shape + (ax_emb.shape[-1],),
          )
          embs.append(ax_emb)
  
      if self._mode == "predict":
          assert self._dropout == 0.0
          emb = jnp.concatenate(embs, -1)
          emb = jnp.reshape(emb, (inputs.shape[0], -1, emb.shape[-1]))
          emb = fastmath.dynamic_slice_in_dim(
              emb, state, inputs.shape[1], axis=1
          )
          self.state = state + inputs.shape[1]
          return inputs + emb
      elif self._dropout == 0:
          # TODO(kitaev): concat-then-reshape (as is the case with dropout enabled)
          # leads to memory blow-up on TPU.
          # emb = jnp.concatenate(embs, -1)
          # return inputs + jnp.reshape(emb, inputs.shape), state
          return inputs + jnp.concatenate(
              [
                  jnp.reshape(emb, inputs.shape[:-1] + (emb.shape[-1],))
                  for emb in embs
              ],
              -1,
          )
      else:
          emb = jnp.concatenate(embs, -1)
          noise_shape = list(emb.shape)
          for dim in self._dropout_broadcast_dims:
              noise_shape[dim] = 1
          keep_prob = 1.0 - self._dropout
          keep = fastmath.random.bernoulli(
              rng, keep_prob, tuple(noise_shape)
          )
          multiplier = keep.astype(inputs.dtype) / keep_prob
          return inputs + jnp.reshape(emb * multiplier, inputs.shape)
  
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 2 for calling the AxialPositionalEncoding obj's underlying member method code with example usage: obj.forward(inputs) and returns: unknown, unknown, unknown