class_docstr: A linear transformation with flexible axes.
  
  Example usage::
  
    >>> import flax.linen as nn
    >>> import jax, jax.numpy as jnp
  
    >>> # equivalent to `nn.Linear(features=4)`
    >>> layer = nn.LinearGeneral(features=4)
    >>> # output features (4, 5)
    >>> layer = nn.LinearGeneral(features=(4, 5))
    >>> params = layer.init(jax.random.key(0), jnp.ones((1, 3)))
    >>> jax.tree_map(jnp.shape, params)
    {'params': {'bias': (4, 5), 'kernel': (3, 4, 5)}}
    >>> # apply transformation on the the second and last axes
    >>> layer = nn.LinearGeneral(features=(4, 5), axis=(1, -1))
    >>> params = layer.init(jax.random.key(0), jnp.ones((1, 3, 6, 7)))
    >>> jax.tree_map(jnp.shape, params)
    {'params': {'bias': (4, 5), 'kernel': (3, 7, 4, 5)}}
  
  Attributes:
    features: int or tuple with number of output features.
    axis: int or tuple with axes to apply the transformation on. For instance,
      (-2, -1) will apply the transformation to the last two axes.
    batch_dims: tuple with batch axes.
    use_bias: whether to add a bias to the output (default: True).
    dtype: the dtype of the computation (default: infer from input and params).
    param_dtype: the dtype passed to parameter initializers (default: float32).
    kernel_init: initializer function for the weight matrix.
    bias_init: initializer function for the bias.
    precision: numerical precision of the computation see `jax.lax.Precision`
      for details.
code: def __init__(
      self,
      in_features: Size | tp.Sequence[Size],
      out_features: Size | tp.Sequence[Size],
      *,
      axis: Axis | tp.Sequence[Axis] = -1,
      batch_axis: tp.Mapping[Axis, Size] = MappingProxyType({}),
      use_bias: bool = True,
      dtype: Dtype | None = None,
      param_dtype: Dtype = jnp.float32,
      kernel_init: initializers.Initializer = default_kernel_init,
      bias_init: initializers.Initializer = initializers.zeros_init(),
      precision: PrecisionLike = None,
      # Deprecated. Will be removed.
      dot_general: DotGeneralT | None = None,
      dot_general_cls: tp.Any = None,
      rngs: rnglib.Rngs,
  ):
      self.in_features = _canonicalize_tuple(in_features)
      self.out_features = _canonicalize_tuple(out_features)
      self.axis = _canonicalize_tuple(axis)
      self.batch_axis = MappingProxyType(batch_axis)
      self.use_bias = use_bias
      self.dtype = dtype
      self.param_dtype = param_dtype
      self.kernel_init = kernel_init
      self.bias_init = bias_init
      self.precision = precision
      self.dot_general = dot_general
      self.dot_general_cls = dot_general_cls
  
      if len(self.in_features) != len(self.axis):
          raise ValueError(
              "in_features and axis must have the same length. "
              f"Got {self.in_features} and {self.axis}."
          )
  
      if batch_axis:
          batch_dims = tuple(batch_axis.keys())
          max_dim = np.max(batch_dims)
          if set(batch_dims) != set(range(max_dim + 1)):
              raise ValueError(
                  "batch_dims %s must be consecutive leading "
                  "dimensions starting from 0." % str(batch_dims)
              )
  
      n_batch_axis = len(self.batch_axis)
      n_in_features = len(self.in_features)
      n_out_features = len(self.out_features)
  
      def kernel_init_wrap(rng, shape, dtype) -> jax.Array:
          flat_shape = (
              np.prod(shape[:n_batch_axis])
              * np.prod(
                  shape[n_batch_axis : n_in_features + n_batch_axis]
              ),
              np.prod(shape[-n_out_features:]),
          )
          flat_shape = jax.tree_map(int, flat_shape)
          kernel = self.kernel_init(rng, flat_shape, dtype)
          if isinstance(kernel, variables.VariableMetadata):
              kernel.value = jnp.reshape(kernel.value, shape)
          else:
              kernel = jnp.reshape(kernel, shape)
  
          return kernel
  
      batch_shape = tuple(self.batch_axis.values())
      kernel_shape = (
          *batch_shape,
          *self.in_features,
          *self.out_features,
      )
      self.kernel = nnx.Param(
          kernel_init_wrap(
              rngs.params(), kernel_shape, self.param_dtype
          )
      )
  
      if self.use_bias:
  
          def bias_init_wrap(rng, shape, dtype) -> jax.Array:
              flat_shape = (int(np.prod(shape)),)
              bias = self.bias_init(rng, flat_shape, dtype)
              if isinstance(bias, variables.VariableMetadata):
                  bias.value = jnp.reshape(bias.value, shape)
              else:
                  bias = jnp.reshape(bias, shape)
              return bias
  
          bias_shape = (*batch_shape, *self.out_features)
          self.bias = nnx.Param(
              bias_init_wrap(
                  rngs.params(), bias_shape, self.param_dtype
              )
          )
      else:
          self.bias = nnx.Param(None)
  
imports: from __future__ import annotations
  import typing
  from types import MappingProxyType
  import jax
  import jax.numpy
  import numpy
  from jax import lax
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import rnglib, variables
  from flax.experimental.nnx.nnx.module import Module
  from flax.experimental.nnx.nnx.nn import dtypes, initializers
label_desc: Use image node_id 1 to create a new LinearGeneral object from inherited base classes: Module with example: obj = LinearGeneral(in_features, out_features)