from collections.abc import Callable
from typing import Any
import jax
import jax.numpy as jnp
import numpy as np
from jax import random
from jax._src import core, dtypes
from .util import gramschmidt_row
[docs]
def zeros(std: Any = 1e-9, dtype: Any = jnp.float_) -> Callable:
"""Build an initializer that initializes tensors with zeros plus small noise.
Examples
--------
>>> import jax, jax.numpy as jnp
>>> from tn4ml.initializers import zeros_init
>>> initializer = zeros_init()
>>> initializer(jax.random.key(42), (2, 2), jnp.float32)
Array([[0., 0.],
[0., 0.]], dtype=float32)
"""
def init(key: Any, shape: core.Shape, dtype: Any = dtype) -> jnp.ndarray:
"""Initialize a tensor.
Parameters
----------
key : Any
Random key.
shape : core.Shape
Shape of the tensor.
dtype : Any
Data type of the tensor.
Returns
-------
jnp.ndarray
Initialized tensor.
"""
return jax.nn.initializers.zeros(key, shape, dtype) + std * random.normal(
key, shape, dtype
)
return init
[docs]
def ones(std: Any = 1e-9, dtype: Any = jnp.float_) -> Callable:
"""Build an initializer that initializes tensors with ones plus small noise.
Examples
--------
>>> import jax, jax.numpy as jnp
>>> from tn4ml.initializers import ones_init
>>> initializer = ones_init()
>>> initializer(jax.random.key(42), (2, 2), jnp.float32)
Array([[1., 1.],
[1., 1.]], dtype=float32)
"""
def init(key: Any, shape: core.Shape, dtype: Any = dtype) -> jnp.ndarray:
"""Initialize a tensor.
Parameters
----------
key : Any
Random key.
shape : core.Shape
Shape of the tensor.
dtype : Any
Data type of the tensor.
Returns
-------
jnp.ndarray
Initialized tensor.
"""
return jax.nn.initializers.ones(key, shape, dtype) + std * random.normal(
key, shape, dtype
)
return init
[docs]
def gramschmidt(dist: str, scale: Any = 1e-2, dtype: Any = jnp.float_) -> Callable:
"""Build an initializer using Gram-Schmidt orthogonalization.
First, arrays are sampled from uniform or normal distribution (specified by `dist` argument).
Parameters
----------
dist : str
Sampling distribution of arrays. Options: `uniform`, `normal`.
scale : Any (Optional). Default = `1e-2`.
Scaling factor for the sampled arrays.
dtype : Any (Optional)
The initializer's default dtype.
Examples
--------
>>> import jax, jax.numpy as jnp
>>> from tn4ml.initializers import gramschmidt_init
>>> initializer = gramschmidt_init('normal')
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)
Array([[ 0.35777482, 0.65598017, 0.6645954 ],
[-0.57674366, -0.40450865, 0.70974606]], dtype=float32)
"""
def init(key: Any, shape: core.Shape, dtype: Any = dtype) -> jnp.ndarray:
"""Initialize a tensor.
Parameters
----------
key : Any
Random key.
shape : core.Shape
Shape of the tensor.
dtype : Any
Data type of the tensor.
Returns
-------
jnp.ndarray
Initialized tensor.
"""
dtype = dtypes.canonicalize_dtype(dtype)
matrix_shape: tuple = (shape[0], int(np.prod(shape[1:])))
if dist == "uniform":
arrays = random.uniform(key, matrix_shape, dtype) * jnp.array(scale, dtype)
elif dist == "normal":
arrays = random.normal(key, matrix_shape, dtype) * jnp.array(scale, dtype)
else:
raise ValueError(
"Sampling only implemented for 'uniform' and 'normal' distributions!"
)
arrays = gramschmidt_row(arrays)
return arrays.reshape(shape)
return init
[docs]
def identity(type: str, std: Any = None, dtype: Any = jnp.float_) -> Callable: # noqa: A002
"""Build an initializer that sets identity values on diagonal or bond dimensions.
Parameters
----------
type : str. Options: 'copy', 'bond'
'copy' = diagonal elements, 'bond' = bond dimension elements
std : Any (Optional)
Additonal noise
dtype : Any (Optional). Default = `jnp.float_`.
The initializer's default dtype.
Examples
--------
>>> import jax, jax.numpy as jnp
>>> from tn4ml.initializers import gramschmidt_init
>>> initializer = identity_init('copy', 1e-2)
>>> initializer(jax.random.key(42), (3, 2), jnp.float32)
Array([[ 1.0061227 , 0.01122588],
[ 0.01137332, 0.99187267],
[-0.00890405, 0.00126231]], dtype=float32)
"""
def init(key: Any, shape: core.Shape, dtype: Any = dtype) -> jnp.ndarray:
"""Initialize a tensor.
Parameters
----------
key : Any
Random key.
shape : core.Shape
Shape of the tensor.
dtype : Any
Data type of the tensor.
Returns
-------
jnp.ndarray
Initialized tensor.
"""
dtype = dtypes.canonicalize_dtype(dtype)
rank = len(shape)
if type == "bond":
tensor = jnp.zeros(shape, dtype=dtype)
if rank == 4:
eye_tensor = jnp.eye(shape[0], shape[1]).reshape(
shape[0], shape[1], 1, 1
)
elif rank == 3:
eye_tensor = jnp.eye(shape[0], shape[1]).reshape(shape[0], shape[1], 1)
else:
raise ValueError("Tensor should have LRP shape")
# Use broadcasting to fill tensor slices
tensor += eye_tensor
elif type == "copy":
# from @joserapa98/tensorkrowch
tensor = jnp.zeros(shape, dtype=dtype)
rank = len(shape)
if rank <= 1:
i: Any = 0
else:
i = np.arange(min(shape), dtype=int)
tensor = tensor.at[(i,) * rank].set(1.0)
else:
raise ValueError(
"Defined only for diagonal and bond dimension identity intialization!"
)
# Add random noise
if std:
tensor += std * random.normal(key, shape, dtype)
return tensor
return init
[docs]
def randn(
std: Any = 1.0,
mean: Any = 0.0,
noise_std: Any = None,
noise_mean: Any = None,
dtype: Any = jnp.float_,
) -> Callable:
"""Build an initializer that samples tensor values from a normal distribution.
Parameters
----------
std : Any (Optional). Default = `1.0`.
Standard deviation of the normal distribution
mean : Any (Optional). Default = `0.0`.
Mean of the normal distribution.
noise_std : Any (Optional). Default = `None`.
The standard deviation of the noise distribution (normal).
noise_mean : Any (Optional). Default = `None`.
The mean of the noise distribution (normal).
dtype : Any (Optional). Default = `jnp.float_`.
The initializer's default dtype.
Examples
--------
>>> import jax, jax.numpy as jnp
>>> from tn4ml.initializers import randn_init
>>> initializer = randn(1e-2)
>>> initializer(jax.random.key(42), (2, 2), jnp.float32)
Array([[ 0.00186935, 0.01065333],
[-0.01559313, -0.01535296]], dtype=float32)
"""
def init(key: Any, shape: core.Shape, dtype: Any = dtype) -> jnp.ndarray:
"""Initialize a tensor.
Parameters
----------
key : Any
Random key.
shape : core.Shape
Shape of the tensor.
dtype : Any
Data type of the tensor.
Returns
-------
jnp.ndarray
Initialized tensor.
"""
dtype = dtypes.canonicalize_dtype(dtype)
tensor = random.normal(key, shape, dtype)
mean_value = 0.0 if mean is None else mean
tensor = mean_value + tensor * std
if noise_std and noise_mean:
noise = random.normal(key, shape, dtype)
tensor += noise_mean + noise * noise_std
return tensor
return init
[docs]
def unitary_matrix(key: Any, shape: core.Shape, dtype: Any = jnp.float_) -> jnp.ndarray:
"""
- from @joserapa98/tensorkrowch.
Generates random unitary matrix from the Haar measure of size n x n.
Unitary matrix is created as described in this `paper
<https://arxiv.org/abs/math-ph/0609050v2>`_.
Parameters
----------
key : Any
Random key.
shape : core.Shape
Shape of the tensor.
dtype : Any
Data type of the tensor.
Returns
-------
jnp.ndarray
Random unitary matrix.
"""
assert shape[0] == shape[1], "Matrix should be square!"
mat = jax.random.normal(key, shape, dtype)
q, r = jnp.linalg.qr(mat)
d = jnp.diagonal(r)
ph = d / jnp.abs(d)
return q @ jnp.diag(ph)
[docs]
def rand_unitary(dtype: Any = jnp.float_) -> Callable:
"""Build an initializer that stacks random unitary matrices.
Parameters
----------
dtype : Any (Optional). Default = `jnp.float_`.
The initializer's default dtype.
Examples
--------
>>> import jax, jax.numpy as jnp
>>> from tn4ml.initializers import rand_unitary
>>> initializer = rand_unitary()
>>> initializer(jax.random.key(42), (2, 2), jnp.float32)
Array([[ 0.11903083, 0.99289054],
[-0.99289054, 0.11903088]], dtype=float32)
>>> tensor = initializer(jax.random.key(42), (2, 2), jnp.float32)
>>> jnp.allclose(tensor @ tensor.T.conj(), jnp.eye(2), atol=1e-6)
True
"""
def init(key: Any, shape: core.Shape, dtype: Any = dtype) -> jnp.ndarray:
"""Initialize a tensor.
Parameters
----------
key : Any
Random key.
shape : core.Shape
Shape of the tensor.
dtype : Any
Data type of the tensor.
Returns
-------
jnp.ndarray
Initialized tensor.
"""
dtype = dtypes.canonicalize_dtype(dtype)
size = max(shape[0], shape[1], shape[2])
size_1 = min(shape[0], size)
size_2 = min(shape[1], size)
if len(shape) == 3:
units = []
for _ in range(shape[2]):
tensor = unitary_matrix(key, (size, size), dtype)
tensor = tensor[:size_1, :size_2]
units.append(tensor)
tensor = jnp.stack(units, axis=-1)
elif len(shape) == 4:
units = []
for _ in range(shape[-2]):
inner_units = []
for _ in range(shape[-1]):
unitary = unitary_matrix(key, (size, size), dtype)
unitary = unitary[:size_1, :size_2]
inner_units.append(unitary)
inner_stack = jnp.stack(inner_units, axis=-1)
units.append(inner_stack)
tensor = jnp.stack(units, axis=-1)
else:
raise ValueError("Only 3 and 4 rank tensors are supported!")
return tensor
return init
