Delay the requirement of the device/pyclient to the first call.
Because jax.jit is annotated on top-level functions, and that backends are initialised only after GoogleInit, we cannot have the client when calling the C++ jit. PiperOrigin-RevId: 328177723 Change-Id: I462385e4a687461bf41c0d3e2875d4736285c42d
Showing
想要评论请 注册 或 登录