oryx.core.state.assign

Assigns a value to a variable.

View aliases

Main aliases

oryx.core.state.module.assign

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).

Args
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'

Returns
The value that was passed in.