class_docstr: Axial positional encoding.
class_docstr_tok: ["Axial","positional","encoding","."]
code: def __init__(
      self,
      shape=(64, 64, 3),
      d_embs=(384, 384, 256),
      kernel_initializer=init.RandomNormalInitializer(1.0),
      dropout=0.0,
      dropout_broadcast_dims=(),
      mode="train",
  ):
      super().__init__()
      self._kernel_initializer = kernel_initializer
      assert len(shape) == len(d_embs)
      self._shape = shape
      self._d_embs = d_embs
  
      if dropout >= 1.0:
          raise ValueError("Dropout rates must be lower than 1.")
      if mode == "train":
          self._dropout = dropout
      else:
          self._dropout = 0.0
      self._dropout_broadcast_dims = dropout_broadcast_dims
      self._mode = mode
  
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 1 to create a new AxialPositionalEncoding object from inherited base classes: layer_base with example: obj = AxialPositionalEncoding(shape, d_embs, kernel_initializer, dropout, dropout_broadcast_dims, mode)