kd.nn.Sequential

kd.nn.Sequential#

class kauldron.modules.Sequential(layers: typing.Sequence[typing.Callable[[...], typing.Any]] = (), inputs: typing.Annotated[typing.Any, <object object at 0x76412092fb90>] = '__KEY_REQUIRED__', outputs: str | None = None, parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)[source]

Bases: flax.linen.module.Module

Like nn.Sequential but allows configuring input and output keys.

inputs

The input key used to pass into the model. E.g. “batch.image”.

Type:

Any

outputs

Optional key to wrap the output in. If set, then the return value is a dictionary {outputs: value}. Otherwise just the value is returned.

Type:

str | None

layers

A sequence of callables (usually nn.Modules) to execute.

Type:

Sequence[Callable[[…], Any]]

layers: Sequence[Callable[..., Any]] = ()
inputs: kontext.Key = '__KEY_REQUIRED__'
outputs: str | None = None
classmethod from_repeated(
layer: flax.linen.module.Module,
repeats: int,
*,
shared_weights: bool = False,
inputs: typing.Annotated[typing.Any,
<object object at 0x76412092fb90>] = '__KEY_REQUIRED__',
outputs: str | None = None,
prefix: typing.Sequence[flax.linen.module.Module] = (),
suffix: typing.Sequence[flax.linen.module.Module] = (),
) → kauldron.modules.models.Sequential[source]

Create a sequential Module by repeating a given layer module.

Effectively constructs a Sequential module with: layers = prefix + [layer] * repeats + suffix

Parameters:

  • layer – The nn.Module to be repeated

  • repeats – How often to repeat the given layer

  • shared_weights – Wether the repeated layers should share weigths.

  • inputs – The input key (see module doc).

  • outputs – Optional key to wrap the output in (see module doc).

  • prefix – A sequence of modules to be added before the repeated layer.

  • suffix – A sequence of modules to be added after the repeated layer.

Returns:

A Sequential Module.

name: str | None = None
parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = None
scope: Scope | None = None