flax.struct 包

flax.struct 包#

用于定义可与 JAX 转换一起使用的自定义类的实用工具。

flax.struct.dataclass(clz=None, **kwargs)[源代码]#

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

注意

请继承 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 – 将被装饰器转换的类。

  • **kwargs – 传递给数据类构造函数的参数。

返回

新类。

class flax.struct.PyTreeNode(*args, **kwargs)[源代码]#

应作为 JAX PyTree 节点的数据类的基类。

关于 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)
目录