variables

class flax.experimental.nnx.BatchStat(value, set_value_hooks=(), get_value_hooks=(), create_value_hooks=(), add_axis_hooks=(), remove_axis_hooks=(), **metadata)

class flax.experimental.nnx.Cache(value, set_value_hooks=(), get_value_hooks=(), create_value_hooks=(), add_axis_hooks=(), remove_axis_hooks=(), **metadata)

class flax.experimental.nnx.Empty

class flax.experimental.nnx.Intermediate(value, set_value_hooks=(), get_value_hooks=(), create_value_hooks=(), add_axis_hooks=(), remove_axis_hooks=(), **metadata)

class flax.experimental.nnx.Param(value, set_value_hooks=(), get_value_hooks=(), create_value_hooks=(), add_axis_hooks=(), remove_axis_hooks=(), **metadata)

class flax.experimental.nnx.Variable(value, set_value_hooks=(), get_value_hooks=(), create_value_hooks=(), add_axis_hooks=(), remove_axis_hooks=(), **metadata)

class flax.experimental.nnx.VariableMetadata(raw_value: 'A', set_value_hooks: 'tuple[SetValueHook[A], ...]' = (), get_value_hooks: 'tuple[GetValueHook[A], ...]' = (), create_value_hooks: 'tuple[CreateValueHook[A], ...]' = (), add_axis_hooks: "tuple[AddAxisHook['Variable[A]'], ...]" = (), remove_axis_hooks: "tuple[RemoveAxisHook['Variable[A]'], ...]" = (), metadata: 'tp.Mapping[str, tp.Any]' = <factory>)

flax.experimental.nnx.with_metadata(initializer, set_value_hooks=(), get_value_hooks=(), create_value_hooks=(), add_axis_hooks=(), remove_axis_hooks=(), **metadata)