netket.models.MPSPeriodic¶
-
class
netket.models.
MPSPeriodic
(hilbert, graph, bond_dim, diag=False, symperiod=None, kernel_init=<function normal.<locals>.init>, dtype=<class 'numpy.complex64'>, parent=<flax.linen.module._Sentinel object>, name=None)[source]¶ Bases:
flax.linen.module.Module
A periodic Matrix Product State (MPS) for a quantum state of discrete degrees of freedom, wrapped as Jax machine.
The MPS is defined as
\[\Psi(s_1,\dots s_N) = \mathrm{Tr} \left[ A[s_1]\dots A[s_N] \right] ,\]for arbitrary local quantum numbers \(s_i\), where \(A[s_1]\) is a matrix of dimension (bdim,bdim), depending on the value of the local quantum number \(s_i\).
-
__init__
(hilbert, graph, bond_dim, diag=False, symperiod=None, kernel_init=<function normal.<locals>.init>, dtype=<class 'numpy.complex64'>, parent=<flax.linen.module._Sentinel object>, name=None)¶ Initialize self. See help(type(self)) for accurate signature.
- Parameters
hilbert (netket.hilbert.abstract_hilbert.AbstractHilbert) –
graph (netket.graph.abstract_graph.AbstractGraph) –
bond_dim (int) –
diag (bool) –
symperiod (bool) –
kernel_init (Callable[[Any, Sequence[int], Any], Any]) –
dtype (Any) –
parent (Optional[Union[Type[Module], Type[Scope], Type[_Sentinel]]]) –
name (str) –
- Return type
- Attributes
-
-
parent
: Optional[Union[Type[Module], Type[Scope], Type[_Sentinel]]] = <flax.linen.module._Sentinel object>¶
-
scope
= None¶
-
- Methods
-
-
apply
(variables, *args, rngs=None, method=None, mutable=False, capture_intermediates=False, **kwargs)¶ Applies a module method to variables and returns output and modified variables.
Note that method should be set if one would like to call apply on a different class method than
__call__
. For instance, suppose a Transformer modules has a method called encode, then the following calls apply on that method:model = Transformer() encoded = model.apply({'params': params}, x, method=Transformer.encode)
If a function instance is provided, the unbound function is used. For instance, the example below is equivalent to the one above:
encoded = model.apply({'params': params}, x, method=model.encode)
Note
method
can also be a function that is not defined inTransformer
. In that case, the function should have at least one argument representing an instance of the Module class:def other_fn(instance, ...): instance.some_module_attr(...) ... model.apply({'params': params}, x, method=other_fn)
- Parameters
variables (
Mapping
[str
,Mapping
[str
,Any
]]) – A dictionary containing variables keyed by variable collections. Seeflax.core.variables
for more details about variables.rngs (
Dict
[str
,Any
]) – a dict of PRNGKeys to initialize the PRNG sequences. The “params” PRNG sequence is used to initialize parameters.method (
Callable
[…,Any
]) – A function to call apply on. This is generally a function in the module. If provided, applies this method. If not provided, applies the__call__
method of the module.mutable (
Union
[bool
,str
,Container
[str
],ForwardRef
]) – Can be bool, str, or list. Specifies which collections should be treated as mutable:bool
: all/no collections are mutable.str
: The name of a single mutable collection.list
: A list of names of mutable collections.capture_intermediates (
Union
[bool
,Callable
[[ForwardRef
,str
],bool
]]) – If True, captures intermediate return values of all Modules inside the “intermediates” collection. By default only the return values of all__call__
methods are stored. A function can be passed to change the filter behavior. The filter function takes the Module instance and method name and returns a bool indicating whether the output of that method invocation should be stored.
- Return type
- Returns
If
mutable
is False, returns output. If any collections are mutable, returns(output, vars)
, wherevars
are is a dict of the modified collections.
-
bind
(variables, *args, rngs=None, mutable=False)¶ Creates an interactive Module instance by binding variables and RNGs.
bind
provides an “interactive” instance of a Module directly without transforming a function withapply
. This is particulary useful for debugging and interactive use cases like notebooks where a function would limit the ability split up code into different cells.Once the variables (and optionally RNGs) are bound to a
Module
it becomes a stateful object. Note that idiomatic JAX is functional and therefore an interactive instance does not mix well well with vanilla JAX APIs.bind()
should only be used for interactive experimentation, and in all other cases we strongly encourage to useapply()
instead.Example:
class AutoEncoder(nn.Module): def setup(self): self.encoder = nn.Dense(3) self.decoder = nn.Dense(5) ae = AutoEncoder() model = ae.bind(variables) z = model.encode(x) x_reconstructed = model.decode(z)
- Parameters
variables (
Mapping
[str
,Mapping
[str
,Any
]]) –- A dictionary containing variables keyed by variable
collections. See
flax.core.variables
for more details about variables.
rngs: a dict of PRNGKeys to initialize the PRNG sequences. mutable: Can be bool, str, or list. Specifies which collections should be
treated as mutable:
bool
: all/no collections are mutable.str
: The name of a single mutable collection.list
: A list of names of mutable collections.rngs (Dict[str, Any]) –
- Returns
A copy of this instance with bound variables and RNGs.
-
clone
(*, parent=None, **updates)¶ Creates a clone of this Module, with optionally updated arguments.
-
get_variable
(col, name, default=None)¶ Retrieves the value of a Variable.
- Parameters
- Return type
~T
- Returns
The value of the input variable, of the default value if the variable doesn’t exist in this scope.
-
has_variable
(col, name)¶ Checks if a variable of given collection and name exists in this Module.
See
flax.core.variables
for more explanation on variables and collections.
-
init
(rngs, *args, method=None, **kwargs)¶ Initializes a module method with variables and returns modified variables.
Jitting init initializes a model lazily using only the shapes of the provided arguments, and avoids computing the forward pass with actual values. Example:
jit_init = jax.jit(SomeModule.init) jit_init(rng, jnp.ones(input_shape, jnp.float32))
- Parameters
- Return type
FrozenDict
[str
,Mapping
[str
,Any
]]- Returns
The initialized variable dict.
-
init_with_output
(rngs, *args, method=None, **kwargs)¶ Initializes a module method with variables and returns output and modified variables.
- Parameters
- Return type
- Returns
(output, vars)`, where
vars
are is a dict of the modified collections.
-
is_mutable_collection
(col)¶ Returns true if the collection col is mutable.
-
kernel_init
(shape, dtype=<class 'jax._src.numpy.lax_numpy.float32'>)¶
-
make_rng
(name)¶ Returns a new RNG key from a given RNG sequence for this Module.
The new RNG key is split from the previous one. Thus, every call to make_rng returns a new RNG key, while still guaranteeing full reproducibility.
TODO: Link to Flax RNG design note.
-
param
(name, init_fn, *init_args)¶ Declares and returns a parameter in this Module.
Parameters are read-only variables in the collection named “params”. See
flax.core.variables
for more details on variables.The first argument of init_fn is assumed to be a PRNG key, which is provided automatically and does not have to be passed using init_args:
mean = self.param('mean', lecun_normal(), (2, 2))
In the example above, the function lecun_normal expects two arguments: key and shape, but only shape has to be provided explicitly; key is set automatically using the PRNG for params that is passed when initializing the module using
init()
.- Parameters
- Return type
~T
- Returns
The value of the initialized parameter.
-
setup
()[source]¶ Initializes a Module lazily (similar to a lazy
__init__
).setup
is called once lazily on a module instance when a module is bound, immediately before any other methods like__call__
are invoked, or before asetup
-defined attribute on self is accessed.This can happen in three cases:
Immediately when invoking
apply()
,init()
orinit_and_output()
.Once the module is given a name by being assigned to an attribute of another module inside the other module’s
setup
method (see__setattr__()
):class MyModule(nn.Module): def setup(self): submodule = Conv(...) # Accessing `submodule` attributes does not yet work here. # The following line invokes `self.__setattr__`, which gives # `submodule` the name "conv1". self.conv1 = submodule # Accessing `submodule` attributes or methods is now safe and # either causes setup() to be called once.
Once a module is constructed inside a method wrapped with
compact()
, immediately before another method is called orsetup
defined attribute is accessed.
-
sow
(col, name, value, reduce_fn=<function <lambda>>, init_fn=<function <lambda>>)¶ Stores a value in a collection.
Collections can be used to collect intermediate values without the overhead of explicitly passing a container through each Module call.
If the target collection is not mutable sow behaves like a no-op and returns False.
Example:
class Foo(nn.Module): @nn.compact def __call__(self, x): h = nn.Dense(4)(x) self.sow('intermediates', 'h', h) return nn.Dense(2)(h) y, state = Foo.apply(params, x, mutable=['intermediates']) print(state['intermediates']) # {'h': (...,)}
By default the values are stored in a tuple and each stored value is appended at the end. This way all intermediates can be tracked when the same module is called multiple times. Alternatively, a custom init/reduce function can be passed:
class Foo(nn.Module): @nn.compact def __call__(self, x): init_fn = lambda: 0 reduce_fn = lambda a, b: a + b self.sow('intermediates', x, h, init_fn=init_fn, reduce_fn=reduce_fn) self.sow('intermediates', x * 2, h, init_fn=init_fn, reduce_fn=reduce_fn) return x y, state = Foo.apply(params, 1, mutable=['intermediates']) print(state['intermediates']) # ==> {'h': 3}
- Parameters
col (
str
) – the variable collection.name (
str
) – the name of the variable.reduce_fn (
Callable
[[~K, ~T], ~K]) – The function used to combine the existing value with the new value the default is to append the value to a tuple.init_fn (
Callable
[[], ~K]) – For the first value stored reduce_fn will be passed the result of init_fn together with the value to be stored. The default is an empty tuple.value (T) –
- Return type
- Returns
True if the value has been stored succesfully, False otherwise.
-
variable
(col, name, init_fn, *init_args)¶ Declares and returns a variable in this Module.
See
flax.core.variables
for more information. See alsoparam()
for a shorthand way to define read-only variables in the “params” collection.Contrary to
param()
, all arguments passing using init_fn should be passed on explictly:key = self.make_rng('stats') mean = self.variable('stats', 'mean', lecun_normal(), key, (2, 2))
In the example above, the function lecun_normal expects two arguments: key and shape, and both have to be passed on. The PRNG for stats has to be provided explicitly when calling
init()
andapply()
.- Parameters
- Return type
Variable
- Returns
A
flax.core.variables.Variable
that can be read or set via “.value” attribute. Throws an error if the variable exists already.
-
-