輔助函數

class flax.nnx.Dict(*args, **kwargs)

class flax.nnx.Sequential(*args, **kwargs)

class flax.nnx.TrainState(graphdef: 'GraphDef[M]', params: 'State', opt_state: 'optax.OptState', step: 'jax.Array', tx: 'optax.GradientTransformation')

replace(**updates)

返回一個新物件，並使用新值取代指定的欄位。