class_docstr: A linear transformation with flexible axes.
  
  Example usage::
  
    >>> import flax.linen as nn
    >>> import jax, jax.numpy as jnp
  
    >>> # equivalent to `nn.Dense(features=4)`
    >>> layer = nn.DenseGeneral(features=4)
    >>> # output features (4, 5)
    >>> layer = nn.DenseGeneral(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.DenseGeneral(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.
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.
      """
      features = _canonicalize_tuple(self.features)
      axis = _canonicalize_tuple(self.axis)
      batch_dims = _canonicalize_tuple(self.batch_dims)
      if batch_dims:
          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)
              )
  
      ndim = inputs.ndim
      n_batch_dims = len(batch_dims)
      axis = _normalize_axes(axis, ndim)
      batch_dims = _normalize_axes(batch_dims, ndim)
      n_axis, n_features = len(axis), len(features)
  
      def kernel_init_wrap(rng, shape, dtype=jnp.float32):
          flat_shape = (
              np.prod(shape[:n_batch_dims])
              * np.prod(shape[n_batch_dims : n_axis + n_batch_dims]),
              np.prod(shape[-n_features:]),
          )
          flat_shape = jax.tree_map(int, flat_shape)
          kernel = self.kernel_init(rng, flat_shape, dtype)
          if isinstance(kernel, meta.AxisMetadata):
              return meta.replace_boxed(
                  kernel, jnp.reshape(kernel.unbox(), shape)
              )
          return jnp.reshape(kernel, shape)
  
      batch_shape = tuple(inputs.shape[ax] for ax in batch_dims)
      # batch and non-contracting dims of input with 1s for batch dims.
      expanded_batch_shape = tuple(
          inputs.shape[ax] if ax in batch_dims else 1
          for ax in range(inputs.ndim)
          if ax not in axis
      )
      kernel_shape = tuple(inputs.shape[ax] for ax in axis) + features
      kernel = self.param(
          "kernel",
          kernel_init_wrap,
          batch_shape + kernel_shape,
          self.param_dtype,
      )
  
      batch_ind = tuple(range(n_batch_dims))
      contract_ind = tuple(range(n_batch_dims, n_axis + n_batch_dims))
  
      if self.use_bias:
  
          def bias_init_wrap(rng, shape, dtype=jnp.float32):
              flat_shape = (
                  np.prod(shape[:n_batch_dims])
                  * np.prod(shape[-n_features:]),
              )
              flat_shape = jax.tree_map(int, flat_shape)
              bias = self.bias_init(rng, flat_shape, dtype)
              if isinstance(bias, meta.AxisMetadata):
                  return meta.replace_boxed(
                      bias, jnp.reshape(bias.unbox(), shape)
                  )
              return jnp.reshape(bias, shape)
  
          bias = self.param(
              "bias",
              bias_init_wrap,
              batch_shape + features,
              self.param_dtype,
          )
      else:
          bias = None
  
      inputs, kernel, bias = 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_dims, batch_ind)),
          precision=self.precision,
      )
      # dot_general output has shape [batch_dims/group_dims] + [feature_dims]
      if self.use_bias:
          # expand bias shape to broadcast bias over batch dims.
          bias = jnp.reshape(bias, expanded_batch_shape + features)
          out += bias
      return out
  
imports: import dataclasses
  from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Union
  import jax
  import jax.numpy
  import numpy
  from jax import eval_shape, lax
  from jax.core import ShapedArray
  from flax.core import meta
  from flax.linen import initializers
  from flax.linen.dtypes import promote_dtype
  from flax.linen.module import Module, compact
label_desc: Use image node_id 1 for calling the DenseGeneral obj's underlying member method code with example usage: obj.__call__(inputs) and returns: out, jnp, meta, jnp, meta