jax.experimental.pallas.run_scoped

jax.experimental.pallas.run_scoped(f, *types, **kw_types)

使用已分配的引用调用函数并返回结果。

位置参数和关键字参数描述了要为每个参数分配的引用类型。除了 jax.experimental.pallas.MemoryRef 之外，每个后端都有自己的一组引用类型。

参数：

• f (Callable[..., Any])

• types (Any)

• kw_types (Any)

返回类型：

Any