class_docstr: A Dense layer with named axes for :meth:`jax.experimental.pjit.pjit`.
  
  .. warning:: This class is hightly EXPERIMENTAL and the API is likely to
      change. For regular (non-pjit) use, please use
      :class:`flax.linen.linear.Dense`.
  
  Attributes:
    features: the number of output features.
    use_bias: whether to add a bias to the output (default: True).
    dtype: the dtype of the computation (default: float32).
    param_dtype: the dtype passed to parameter initializers (default: float32).
    precision: numerical precision of the computation see `jax.lax.Precision`
      for details.
    kernel_init: initializer function for the weight matrix.
    bias_init: initializer function for the bias.
docstr: Applies a linear transformation to the inputs along the last dimension.
  
  Args:
    inputs: The nd-array to be transformed.
  
  Returns:
    The transformed input.
docstr_tok: ["Applies","a","linear","transformation","to","the","inputs","along","the","last","dimension",".","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 the last dimension.
  
      Args:
        inputs: The nd-array to be transformed.
  
      Returns:
        The transformed input.
      """
      inputs = jnp.asarray(inputs, self.dtype)
      kernel = param_with_axes(
          "kernel",
          self.kernel_init,
          (inputs.shape[-1], self.features),
          self.param_dtype,
          axes=self.kernel_axes,
      )
      kernel = jnp.asarray(kernel, 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
      y = dot_general(
          inputs,
          kernel,
          (((inputs.ndim - 1,), (0,)), ((), ())),
          precision=self.precision,
      )
      if self.use_bias:
          bias = param_with_axes(
              "bias",
              self.bias_init,
              (self.features,),
              self.param_dtype,
              axes=(self.kernel_axes[-1],),
          )
          bias = jnp.asarray(bias, self.dtype)
          y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))
      return y
  
imports: import dataclasses
  from typing import Any, Callable, Iterable, Optional, Tuple, Union
  import jax.numpy
  from jax import lax
  from flax import linen
  from flax.linen import initializers
  from flax.linen.linear import DotGeneralT, PrecisionLike
  from flax.linen.partitioning import param_with_axes, with_sharding_constraint
label_desc: Use image node_id 1 for calling the Dense obj's underlying member method code with example usage: obj.__call__(inputs) and returns: y