parameters: (sampler, model, variables, state, chain_length)
return_statement: return σ, new_state
docstring: Internal method used for jitting calls.
docstring_summary: Internal method used for jitting calls.
docstring_tokens: Internal method used for jitting calls .
function: def _sample_chain(sampler, model, variables, state, chain_length):
      """
      Internal method used for jitting calls.
      """
      if "cache" in variables:
          variables, _ = variables.pop("cache")
  
      def scan_fun(carry, index):
          σ, cache, key = carry
          if cache:
              _variables = {**variables, "cache": cache}
          else:
              _variables = variables
          new_key, key = jax.random.split(key)
  
          p, mutables = model.apply(
              _variables,
              σ,
              index,
              method=model._conditional,
              mutable=["cache"],
          )
          if "cache" in mutables:
              cache = mutables["cache"]
          else:
              cache = None
  
          local_states = jnp.asarray(sampler.hilbert.local_states, dtype=sampler.dtype)
          new_σ = batch_choice(key, local_states, p)
          σ = σ.at[:, index].set(new_σ)
  
          return (σ, cache, new_key), None
  
      new_key, key_init, key_scan = jax.random.split(state.key, 3)
  
      # We just need a buffer for `σ` before generating each sample
      # The result does not depend on the initial contents in it
      σ = jnp.zeros(
          (chain_length * sampler.n_chains_per_rank, sampler.hilbert.size),
          dtype=sampler.dtype,
      )
  
      # Initialize `cache` before generating each sample,
      # even if `variables` is not changed and `reset` is not called
      cache = sampler._init_cache(model, σ, key_init)
  
      indices = jnp.arange(sampler.hilbert.size)
      (σ, _, _), _ = jax.lax.scan(scan_fun, (σ, cache, key_scan), indices)
      σ = σ.reshape((chain_length, sampler.n_chains_per_rank, sampler.hilbert.size))
  
      new_state = state.replace(key=new_key)
      return σ, new_state
function_tokens: def _sample_chain ( sampler , model , variables , state , chain_length ) : if "cache" in variables : variables , _ = variables . pop ( "cache" ) def scan_fun ( carry , index ) : σ, ache, ey  arry if cache : _variables = { * * variables , "cache" : cache } else : _variables = variables new_key , key = jax . random . split ( key ) p , mutables = model . apply ( _variables , σ, index , method = model . _conditional , mutable = [ "cache" ] , ) if "cache" in mutables : cache = mutables [ "cache" ] else : cache = None local_states = jnp . asarray ( sampler . hilbert . local_states , dtype = sampler . dtype ) new_σ  atch_choice( k ey, ocal_states, ) σ  .a t [: , i dex]. s e t(n e w_σ) return ( σ, ache, ew_key) , one new_key , key_init , key_scan = jax . random . split ( state . key , 3 ) # We just need a buffer for `σ` before generating each sample # The result does not depend on the initial contents in it σ  np. z eros( ( chain_length * sampler . n_chains_per_rank , sampler . hilbert . size ) , dtype = sampler . dtype , ) # Initialize `cache` before generating each sample, # even if `variables` is not changed and `reset` is not called cache = sampler . _init_cache ( model , σ, ey_init) indices = jnp . arange ( sampler . hilbert . size ) ( σ, , ) , ax. l ax. s can( s can_fun, σ ,  c che,  k y_scan), i dices) σ  .r e shape(( c h ain_length,  s mpler.n _ chains_per_rank,  s mpler.h i lbert.s i ze)) new_state = state . replace ( key = new_key ) return σ, ew_state