flax.struct 包

flax.struct 包#

用于定义可与 jax 变换一起使用的自定义类的实用程序。

flax.struct.dataclass(clz, **kwargs)[source]#

创建一个可以传递给函数式变换的类。

注意

继承 PyTreeNode 而不是为了避免在使用 PyType 时出现类型检查问题。

Jax 变换（如 jax.jit 和 jax.grad）需要不可变的对象，可以使用 jax.tree_util 方法对其进行映射。 dataclass 装饰器使定义可以安全地传递给 Jax 的自定义类变得容易。 例如

>>> from flax import struct
>>> import jax
>>> from typing import Any, Callable

>>> @struct.dataclass
... class Model:
...   params: Any
...   # use pytree_node=False to indicate an attribute should not be touched
...   # by Jax transformations.
...   apply_fn: Callable = struct.field(pytree_node=False)

...   def __apply__(self, *args):
...     return self.apply_fn(*args)

>>> params = {}
>>> params_b = {}
>>> apply_fn = lambda v, x: x
>>> model = Model(params, apply_fn)

>>> # model.params = params_b  # Model is immutable. This will raise an error.
>>> model_b = model.replace(params=params_b)  # Use the replace method instead.

>>> # This class can now be used safely in Jax to compute gradients w.r.t. the
>>> # parameters.
>>> model = Model(params, apply_fn)
>>> loss_fn = lambda model: 3.
>>> model_grad = jax.grad(loss_fn)(model)

请注意，数据类具有自动生成的 __init__，其中构造函数的参数和创建的实例的属性一一对应。这种对应关系使这些对象成为有效的容器，可以与 JAX 变换和更广泛地与 jax.tree_util 库一起使用。

有时需要“智能构造函数”，例如，因为某些属性可以（可选地）从其他属性派生。在 Flax 数据类中执行此操作的方法是创建提供智能构造函数的静态方法或类方法。这样一来，jax.tree_util 使用的简单构造函数将保留。考虑以下示例

>>> @struct.dataclass
... class DirectionAndScaleKernel:
...   direction: jax.Array
...   scale: jax.Array

...   @classmethod
...   def create(cls, kernel):
...     scale = jax.numpy.linalg.norm(kernel, axis=0, keepdims=True)
...     direction = direction / scale
...     return cls(direction, scale)
参数

clz – 将被装饰器转换的类。

返回值

新类。

class flax.struct.PyTreeNode(*args, **kwargs)[source]#

应该充当 JAX pytree 节点的 dataclass 的基类。

有关 jax.tree_util 行为，请参见 flax.struct.dataclass。此基类还避免了在使用 PyType 时出现类型检查错误。

示例

>>> from flax import struct
>>> import jax
>>> from typing import Any, Callable

>>> class Model(struct.PyTreeNode):
...   params: Any
...   # use pytree_node=False to indicate an attribute should not be touched
...   # by Jax transformations.
...   apply_fn: Callable = struct.field(pytree_node=False)

...   def __apply__(self, *args):
...     return self.apply_fn(*args)

>>> params = {}
>>> params_b = {}
>>> apply_fn = lambda v, x: x
>>> model = Model(params, apply_fn)

>>> # model.params = params_b  # Model is immutable. This will raise an error.
>>> model_b = model.replace(params=params_b)  # Use the replace method instead.

>>> # This class can now be used safely in Jax to compute gradients w.r.t. the
>>> # parameters.
>>> model = Model(params, apply_fn)
>>> loss_fn = lambda model: 3.
>>> model_grad = jax.grad(loss_fn)(model)
内容