Ejemplos de global_state en Python

Ejemplo n.º 1

def _update_global_jit_state(**kw):
  gs = jax_jit.global_state()
  context = gs.extra_jit_context or _GlobalExtraJitContext()
  gs.extra_jit_context = context._replace(**kw)

Ejemplo n.º 2

def update_global_jit_state(**kw):
    gs = jax_jit.global_state()
    context = gs.extra_jit_context or GlobalJitState()
    gs.extra_jit_context = context._replace(**kw)