(cls, key, in_dim, out_dim, dtype, *args, **kwargs)

(cls, name, common_params, *args, **kwargs)

(key, shape, dtype)

(**kwargs)

(*a)

(cls, common_params, x, *args, **kwargs)

(params, frozen_params, thread_ids)

(x, w, eps=1e-5)

(x)

(frozen_noiser_params, noiser_params, param, Z, pop_size)

(frozen_noiser_params, iterinfo, param, key)

(BIG_RAND_MATRIX, frozen_noiser_params, iterinfo, param, key)

(cls, frozen_noiser_params, noiser_params, param, base_key, fitnesses, iterinfos, map_classification)

(x, params)

()

(cls, frozen_noiser_params, noiser_params, raw_scores, num_episodes_list=None)

(cls, params, config)

(cls, frozen_noiser_params, noiser_params, param, base_key, iterinfo, x)

(cls, frozen_noiser_params, noiser_params, param, base_key, iterinfo, x)

(cls, frozen_noiser_params, noiser_params, params, base_keys, fitnesses, iterinfos, es_map)

(base_key_int32, new_int32)

(params, tokens)

(batch_size, population_size)

(LOG2TABLE, logits, target)

(cls, frozen_noiser_params, noiser_params, param, base_key, iterinfo)

(BIG_RAND_MATRIX, frozen_noiser_params, iterinfo, param, key)

(path: str | Path)

(common, **kwargs)

(key, n_embd, n_layer, vocab_size, ctx_len, dtype=jnp.float32)

()

(model: any, path: str | Path, overwrite: bool = False)

(model_cls, key, n_layer, *args, **kwargs)

(x)

(params, base_key, scan_map)

(x)

(cls, common_params, *args, **kwargs)

(cls, common_params, *args, **kwargs)

(cls, common_params, x, *args, **kwargs)

(cls, common_params, x, *args, **kwargs)

(cls, common_params, x, *args, **kwargs)

(cls, common_params, x)

(cls, common_params, x, state, *args, **kwargs)

(cls, common_params, x, state, *args, **kwargs)

(cls, common_params, x, state, *args, **kwargs)

(noiser_params, base_sigma, ppf, param, key, scores, iterinfo, frozen_noiser_params)

(frozen_noiser_params, noiser_params, param, key, scores, iterinfo)

(frozen_noiser_params, noiser_params, param, key, scores, iterinfo)

(params, in_tokens, out_tokens)

(x, block)

(x, inputs)

(cls, frozen_noiser_params, noiser_params, param, base_key, iterinfo, x)

(params, optimizer, input_tokens, output_tokens)

(params, thread_ids)

(frozen_noiser_params, frozen_params, es_tree_key, LOG2TABLE, noiser_params, params, input_tokens, target_toke

(state, inputs)

(params, in_tokens, out_tokens)

(x, params)

(cls, key, shape, scale, raw_value, dtype, *args, **kwargs)

(cls, key, vocab_size, hidden_dim, dtype, *args, **kwargs)

(cls, key, in_dim, out_dim, use_bias, dtype, *args, **kwargs)

(cls, key, in_dim, out_dim, hidden_dims, use_bias, activation, dtype, *args, **kwargs)

(cls, key, hidden_dim, dtype, *args, **kwargs)

(cls, key, hidden_dim, dtype, *args, **kwargs)

(cls, key, hidden_dim, dtype, *args, **kwargs)

(cls, key, vocab_size, n_layer, hidden_dim, dtype, *args, **kwargs)

(param, base_key, scan_tuple)

(Z, inputs)