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.
docstr: Applies a linear transformation to the inputs along multiple dimensions.
  
  Args:
    inputs: The nd-array to be transformed.
  
  Returns:
    The transformed input.
docstr_tok: ["Applies","a","linear","transformation","to","the","inputs","along","multiple","dimensions",".","Args",":","inputs",":","The","nd-array","to","be","transformed",".","Returns",":","The","transformed","input","."]
code: def __call__(self, inputs: Array) -> Array:
      """Applies a linear transformation to the inputs along multiple dimensions.
  
      Args:
        inputs: The nd-array to be transformed.
  
      Returns:
        The transformed input.
      """
  
      ndim = inputs.ndim
      n_batch_dims = len(self.batch_axis)
      axis = _normalize_axes(self.axis, ndim)
      batch_axis = _normalize_axes(tuple(self.batch_axis.keys()), ndim)
      n_axis = len(axis)
  
      # batch and non-contracting dims of input with 1s for batch dims.
      expanded_batch_shape = tuple(
          inputs.shape[ax] if ax in batch_axis else 1
          for ax in range(inputs.ndim)
          if ax not in axis
      )
      kernel = self.kernel
      bias = self.bias
  
      batch_ind = tuple(range(n_batch_dims))
      contract_ind = tuple(range(n_batch_dims, n_axis + n_batch_dims))
  
      inputs, kernel, bias = dtypes.promote_dtype(
          inputs, kernel, bias, dtype=self.dtype
      )
  
      if self.dot_general_cls is not None:
          dot_general = self.dot_general_cls()
      elif self.dot_general is not None:
          dot_general = self.dot_general
      else:
          dot_general = lax.dot_general
      out = dot_general(
          inputs,
          kernel,
          ((axis, contract_ind), (batch_axis, batch_ind)),
          precision=self.precision,
      )
      # dot_general output has shape [batch_dims/group_dims] + [feature_dims]
      if bias is not None:
          # expand bias shape to broadcast bias over batch dims.
          bias = jnp.reshape(
              bias, (*expanded_batch_shape, *self.out_features)
          )
          out += bias
      return out
  
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 2 for calling the LinearGeneral obj's underlying member method code with example usage: obj.__call__(inputs) and returns: out