class_docstr: Affine transformation layer.
  
  Described in:
  Touvron et al, "ResMLP: Feedforward networks for image classification
  with data-efficient training", 2021.
  
  Performs an affine transformation on the final dimension of the input tensor.
class_docstr_tok: ["Affine","transformation","layer",".","Described","in",":","Touvron","et","al",",","``","ResMLP",":","Feedforward","networks","for","image","classification","with","data-efficient","training","''",",","2021",".","Performs","an","affine","transformation","on","the","final","dimension","of","the","input","tensor","."]
code: def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      n = x.shape[-1]
      scale = self.param("scale", self.scale_init, (n,))
      if self.use_bias:
          bias = self.param("bias", self.bias_init, (n,))
      else:
          bias = 0.0
      return scale * x + bias
  
imports: from typing import Callable, Iterable, Optional, Sequence
  import flax.linen
  import jax
  from jax.nn import initializers
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the Affine obj's underlying member method code with example usage: obj.__call__(x) and returns: unknown