| """Hackable surface for the multiscale polynomial memory research code.""" |
| from __future__ import annotations |
|
|
| from .integration import IntegrationResult, integrate_system |
| from .models import AssocMemHiPPO, SalienceHiPPO, VanillaHiPPO, legendre_orthonormal_basis01 |
| from .polynomials import ( |
| EvalFn, |
| SystemParams, |
| fourier_val, |
| get_L_vector, |
| get_output_vector, |
| get_system_params, |
| legval, |
| legsval, |
| legval_jax, |
| legsval_jax, |
| make_m_mat, |
| ) |
| from .synthetic import ( |
| get_ou_sample, |
| oracle_mse_censored_task, |
| oracle_info_gain, |
| oracle_mse_rbf_periodic_censor, |
| oracle_info_gain_rbf, |
| make_copying_task, |
| make_censored_task, |
| make_induction_head_multioutput_s2s_task, |
| make_induction_head_seq_to_seq_task, |
| make_induction_head_task, |
| make_implicit_measure_task, |
| make_multiplexing_task, |
| whitesignal, |
| wray_and_green_output, |
| make_simple_repetition_task |
| ) |
| from .utils import time_block, unvec, vec |
|
|
| __all__ = [ |
| "AssocMemHiPPO", |
| "EvalFn", |
| "IntegrationResult", |
| "SystemParams", |
| "fourier_val", |
| "get_L_vector", |
| "get_output_vector", |
| "get_ou_sample", |
| "get_system_params", |
| "SalienceHiPPO", |
| "VanillaHiPPO", |
| "integrate_system", |
| "legendre_orthonormal_basis01", |
| "legval", |
| "legsval", |
| "make_induction_head_multioutput_s2s_task", |
| "make_induction_head_seq_to_seq_task", |
| "make_induction_head_task", |
| "make_implicit_measure_task", |
| "make_m_mat", |
| "make_multiplexing_task", |
| "time_block", |
| "unvec", |
| "vec", |
| ] |
|
|
