jax.named_scope

jax.named_scope(name)

一个上下文管理器，将用户指定的名称添加到 JAX 名称堆栈中。

当为了即时编译到 XLA（或其他后端，如 TensorFlow）而暂存计算时，默认情况下，JAX 不会保留它遇到的 Python 函数的名称（或其他源元数据）。这会使调试程序的暂存（和/或编译）表示变得复杂，因为每个正在执行的操作的上下文信息有限。

named_scope 告诉 JAX 暂存给定的函数，并在底层操作上添加额外的注解。JAX 在内部使用名称堆栈来跟踪这些注解。当暂存的程序使用 XLA 编译时，这些注解会被保留，并会出现在调试工具中，例如 TensorBoard 中的 TensorFlow Profiler。当使用 experimental.jax2tf.convert() 将 JAX 程序暂存到 TensorFlow 时，也会保留名称。

参数:

name (str) – 用于命名在名称作用域内创建的所有操作的前缀。

产生:

产生 None，但进入一个上下文中，其中 name 将被附加到活动名称堆栈。

返回类型:

Generator[None, None, None]

示例

named_scope 可以用作编译函数内的上下文管理器

>>> import jax
>>>
>>> @jax.jit
... def layer(w, x):
...   with jax.named_scope("dot_product"):
...     logits = w.dot(x)
...   with jax.named_scope("activation"):
...     return jax.nn.relu(logits)

它也可以用作装饰器

>>> @jax.jit
... @jax.named_scope("layer")
... def layer(w, x):
...   logits = w.dot(x)
...   return jax.nn.relu(logits)