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