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