flax.linen.Dense

class flax.linen.Dense(features, use_bias=True, dtype=None, param_dtype=<class 'jax.numpy.float32'>, precision=None, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, parent=<flax.linen.module._Sentinel object>, name=None)[source]

A linear transformation applied over the last dimension of the input.

Parameters
  • features (int) –

  • use_bias (bool) –

  • dtype (Optional[Any]) –

  • param_dtype (Any) –

  • precision (Union[None, str, jax._src.lax.lax.Precision, Tuple[str, str], Tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]]) –

  • kernel_init (Callable[[Any, Tuple[int, ...], Any], Any]) –

  • bias_init (Callable[[Any, Tuple[int, ...], Any], Any]) –

  • parent (Optional[Union[Type[flax.linen.module.Module], Type[flax.core.scope.Scope], Type[flax.linen.module._Sentinel]]]) –

  • name (str) –

Return type

None

features

the number of output features.

Type

int

use_bias

whether to add a bias to the output (default: True).

Type

bool

dtype

the dtype of the computation (default: infer from input and params).

Type

Optional[Any]

param_dtype

the dtype passed to parameter initializers (default: float32).

Type

Any

precision

numerical precision of the computation see jax.lax.Precision for details.

Type

Union[None, str, jax._src.lax.lax.Precision, Tuple[str, str], Tuple[jax._src.lax.lax.Precision, jax._src.lax.lax.Precision]]

kernel_init

initializer function for the weight matrix.

Type

Callable[[Any, Tuple[int, …], Any], Any]

bias_init

initializer function for the bias.

Type

Callable[[Any, Tuple[int, …], Any], Any]

__call__(inputs)[source]

Applies a linear transformation to the inputs along the last dimension.

Parameters

inputs (Any) – The nd-array to be transformed.

Returns

The transformed input.

Return type

Any

Methods

bias_init(shape[, dtype])

An initializer that returns a constant array full of zeros.

kernel_init(shape[, dtype])