oryx.core.state.assign

Assigns a value to a variable.

oryx.core.state.assign(
    value, *, name: str, key=None, mode: str = 'clobber'
)

In a stateful function, assign is used define state updates. In particular, when a function with an assign is transformed using init, it returns a Module whose call_and_update returns the values tagged as assign as its second output. init requires that an assigned value must have a matching variable (as defined by the name).

value JAX value to be assigned.
name string name for the value.
key JAX value that is used to tie in value. Default value: None
mode string name for sow mode (see harvest documentation). Default value: 'clobber'

The value that was passed in.