class_docstr: 2-Dimensional Transformer encoder block with Axial attention.
  
  This block is similar to Encoder1DBlock, where instead of `self-attention+MLP`
  we have `row-self-attention + col-self-attention + MLP`.
  
  Attributes:
    mlp_dim: dimension of the mlp on top of attention block.
    attention_configs: Configs pass to the self-attention func.
    attention_fn: Type of the sel-attention function.
    dropout_rate: Dropout used in the mlp block.
    attention_dropout_rate: Dropout for attention heads.
    factorization_axis: Axis over which we run attention.
    post_sa_fn: Function to be applied on the output of self-attention block.
    dtype: The dtype of the computation.
class_docstr_tok: ["2-Dimensional","Transformer","encoder","block","with","Axial","attention",".","This","block","is","similar","to","Encoder1DBlock",",","where","instead","of","`","self-attention+MLP","`","we","have","`","row-self-attention","+","col-self-attention","+","MLP","`",".","Attributes",":","mlp_dim",":","dimension","of","the","mlp","on","top","of","attention","block",".","attention_configs",":","Configs","pass","to","the","self-attention","func",".","attention_fn",":","Type","of","the","sel-attention","function",".","dropout_rate",":","Dropout","used","in","the","mlp","block",".","attention_dropout_rate",":","Dropout","for","attention","heads",".","factorization_axis",":","Axis","over","which","we","run","attention",".","post_sa_fn",":","Function","to","be","applied","on","the","output","of","self-attention","block",".","dtype",":","The","dtype","of","the","computation","."]
code: def get_drop_pattern(self, x, deterministic):
      if not deterministic and self.droplayer_p:
          shape = (x.shape[0],) + (1,) * (x.ndim - 1)
          return jax.random.bernoulli(
              self.make_rng("dropout"), self.droplayer_p, shape
          ).astype("float32")
      else:
          return 0.0
  
imports: import abc
  import enum
  import functools
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Iterable
  from absl import logging
  from flax import linen
  import jax
  from jax import lax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import attention_layers
label_desc: Use image node_id 1 for calling the EncoderAxialBlock obj's underlying member method code with example usage: obj.get_drop_pattern(x, deterministic) and returns: jax, int