{ "schema_version": 1, "source": { "repository": "https://github.com/Axym-Labs/axosim", "commit": "306a51ed950b411e8858af622d48062b28e4fbfe" }, "exports": [ { "name": "AXOSIM_POPULATION_PROFILE", "module": "simulation_contract", "implementation": "AXOSIM_POPULATION_PROFILE", "page": "api/simulation-contract" }, { "name": "MILLION_NEURON_REALTIME_CONTRACT", "module": "simulation_contract", "implementation": "MILLION_NEURON_REALTIME_CONTRACT", "page": "api/simulation-contract" }, { "name": "AxoMamba", "module": "axomamba", "implementation": "AxoMamba", "page": "api/mamba" }, { "name": "AxoMambaConfig", "module": "axomamba", "implementation": "AxoMambaConfig", "page": "api/mamba" }, { "name": "AxoSimGRU", "module": "block_forecast", "implementation": "CausalBlockForecastModel", "page": "api/block-forecast" }, { "name": "AxoSimLite", "module": "support_surrogate", "implementation": "AdaptiveSupportP4Surrogate", "page": "api/lite" }, { "name": "AxoSimMamba", "module": "axomamba", "implementation": "AxoMamba", "page": "api/mamba" }, { "name": "AxoSimModelProfile", "module": "model_family", "implementation": "AxoSimModelProfile", "page": "api/model-family" }, { "name": "AxoSimPopulation", "module": "interfaces", "implementation": "AxoSimPopulation", "page": "api/interfaces" }, { "name": "CudaGraphAdaptationStep", "module": "adaptation", "implementation": "CudaGraphAdaptationStep", "page": "api/adaptation" }, { "name": "HomeostaticThresholdController", "module": "activity", "implementation": "HomeostaticThresholdController", "page": "api/activity" }, { "name": "AXOSIM_MODEL_PROFILES", "module": "model_family", "implementation": "AXOSIM_MODEL_PROFILES", "page": "api/model-family" }, { "name": "BranchELM", "module": "model", "implementation": "BranchELM", "page": "api/model" }, { "name": "BranchELMConfig", "module": "model", "implementation": "BranchELMConfig", "page": "api/model" }, { "name": "LargePopulationSimulationContract", "module": "simulation_contract", "implementation": "LargePopulationSimulationContract", "page": "api/simulation-contract" }, { "name": "MixedQuantizedNeuronBehaviorBank", "module": "population", "implementation": "MixedQuantizedNeuronBehaviorBank", "page": "api/population" }, { "name": "PopulationInferenceProfile", "module": "simulation_contract", "implementation": "PopulationInferenceProfile", "page": "api/simulation-contract" }, { "name": "ProceduralMorphologyConnectomeRouter", "module": "connectome", "implementation": "ProceduralMorphologyConnectomeRouter", "page": "api/connectome" }, { "name": "QuantizedSynapticEfficacyBank", "module": "synapse", "implementation": "QuantizedSynapticEfficacyBank", "page": "api/synapse" }, { "name": "compact_active_sources", "module": "connectome", "implementation": "compact_active_sources", "page": "api/connectome" }, { "name": "dequantize_synaptic_efficacies", "module": "synapse", "implementation": "dequantize_synaptic_efficacies", "page": "api/synapse" }, { "name": "dequantize_mixed_neuron_behavior_parameters", "module": "population", "implementation": "dequantize_mixed_neuron_behavior_parameters", "page": "api/population" }, { "name": "quantize_synaptic_efficacies", "module": "synapse", "implementation": "quantize_synaptic_efficacies", "page": "api/synapse" }, { "name": "quantize_mixed_neuron_behavior_parameters", "module": "population", "implementation": "quantize_mixed_neuron_behavior_parameters", "page": "api/population" }, { "name": "create_axosim_profile", "module": "model_family", "implementation": "create_axosim_profile", "page": "api/model-family" } ], "pages": [ { "slug": "api/index", "modules": [ "__init__" ] }, { "slug": "api/interfaces", "modules": [ "interfaces" ] }, { "slug": "api/model-family", "modules": [ "model_family" ] }, { "slug": "api/mamba", "modules": [ "axomamba", "mamba_official" ] }, { "slug": "api/temporal-core", "modules": [ "temporal_core" ] }, { "slug": "api/block-forecast", "modules": [ "block_forecast" ] }, { "slug": "api/lite", "modules": [ "support_surrogate" ] }, { "slug": "api/model", "modules": [ "model" ] }, { "slug": "api/checkpoint", "modules": [ "checkpoint" ] }, { "slug": "api/adaptation", "modules": [ "adaptation" ] }, { "slug": "api/simulation-contract", "modules": [ "simulation_contract" ] }, { "slug": "api/population", "modules": [ "population", "deployment" ] }, { "slug": "api/synapse", "modules": [ "synapse" ] }, { "slug": "api/connectome", "modules": [ "connectome" ] }, { "slug": "api/activity", "modules": [ "activity" ] }, { "slug": "api/data", "modules": [ "data" ] }, { "slug": "api/neuronio", "modules": [ "neuronio_raw" ] }, { "slug": "api/metrics", "modules": [ "metrics", "evaluate" ] }, { "slug": "api/training", "modules": [ "train" ] }, { "slug": "api/inference-benchmark", "modules": [ "inference_benchmark" ] }, { "slug": "api/axobench", "modules": [ "axobench_iteration" ] }, { "slug": "api/setup", "modules": [ "setup_workflow" ] }, { "slug": "api/cli", "modules": [ "cli" ] }, { "slug": "api/setup-cli", "modules": [ "setup_workflow" ] }, { "slug": "api/axobench-cli", "modules": [ "axobench_iteration" ] } ], "symbols": [ { "module": "interfaces", "name": "AxoSimPopulation", "kind": "class", "page": "api/interfaces", "line": 24, "end_line": 470, "signature": "AxoSimPopulation(self, neuron: AxoSimLite, *, morphology_indices: torch.Tensor, contact_branch_indices: torch.Tensor) -> None", "parameters": [ { "name": "neuron", "type": "AxoSimLite", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "contact_branch_indices", "type": "torch.Tensor", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, neuron: AxoSimLite, *, morphology_indices: torch.Tensor, contact_branch_indices: torch.Tensor) -> None", "parameters": [ { "name": "neuron", "type": "AxoSimLite", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "contact_branch_indices", "type": "torch.Tensor", "default": null, "kind": "keyword_only" } ], "line": 27, "end_line": 76, "return_type": "None", "decorators": [] }, { "name": "population_size", "kind": "property", "signature": "AxoSimPopulation.population_size: int", "parameters": [], "line": 79, "end_line": 80, "return_type": "int", "decorators": [ "property" ] }, { "name": "runtime_coefficients_per_neuron", "kind": "property", "signature": "AxoSimPopulation.runtime_coefficients_per_neuron: int", "parameters": [], "line": 83, "end_line": 84, "return_type": "int", "decorators": [ "property" ] }, { "name": "synaptic_efficacies_per_neuron", "kind": "property", "signature": "AxoSimPopulation.synaptic_efficacies_per_neuron: int", "parameters": [], "line": 87, "end_line": 88, "return_type": "int", "decorators": [ "property" ] }, { "name": "adaptation_parameters", "kind": "method", "signature": "adaptation_parameters(self) -> Iterator[nn.Parameter]", "parameters": [], "line": 90, "end_line": 93, "return_type": "Iterator[nn.Parameter]", "decorators": [] }, { "name": "trainable_adaptation_parameters", "kind": "method", "signature": "trainable_adaptation_parameters(self) -> Iterator[nn.Parameter]", "parameters": [], "line": 95, "end_line": 100, "return_type": "Iterator[nn.Parameter]", "decorators": [] }, { "name": "enable_adaptation_training", "kind": "method", "signature": "enable_adaptation_training(self, *, behavior: bool=True, morphology: bool=True, synaptic: bool=True) -> None", "parameters": [ { "name": "behavior", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "morphology", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "synaptic", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "line": 102, "end_line": 113, "return_type": "None", "decorators": [] }, { "name": "enable_full_training", "kind": "method", "signature": "enable_full_training(self) -> None", "parameters": [], "line": 115, "end_line": 117, "return_type": "None", "decorators": [] }, { "name": "synaptic_efficacies", "kind": "method", "signature": "synaptic_efficacies(self, synaptic_log_efficacy: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "synaptic_log_efficacy", "type": "torch.Tensor | None", "default": "None", "kind": "positional_or_keyword", "resolved_default": null } ], "line": 119, "end_line": 134, "return_type": "torch.Tensor", "decorators": [] }, { "name": "forward", "kind": "method", "signature": "forward(self, contact_inputs: torch.Tensor, *, synaptic_log_efficacy: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "contact_inputs", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "synaptic_log_efficacy", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 198, "end_line": 211, "return_type": "torch.Tensor", "decorators": [] }, { "name": "forward_sparse_contacts", "kind": "method", "signature": "forward_sparse_contacts(self, event_summary_indices: torch.Tensor, event_contact_indices: torch.Tensor, event_values: torch.Tensor, *, time_steps: int, synaptic_log_efficacy: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "event_summary_indices", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "event_contact_indices", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "event_values", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "time_steps", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "synaptic_log_efficacy", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 213, "end_line": 327, "return_type": "torch.Tensor", "decorators": [] }, { "name": "forward_batch", "kind": "method", "signature": "forward_batch(self, contact_inputs: torch.Tensor, *, synaptic_log_efficacy: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "contact_inputs", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "synaptic_log_efficacy", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 329, "end_line": 392, "return_type": "torch.Tensor", "decorators": [] }, { "name": "forward_tokens", "kind": "method", "signature": "forward_tokens(self, tokens: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "tokens", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 394, "end_line": 418, "return_type": "torch.Tensor", "decorators": [] }, { "name": "forward_token_batch", "kind": "method", "signature": "forward_token_batch(self, tokens: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "tokens", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 420, "end_line": 470, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [ "nn.Module" ] }, { "module": "model_family", "name": "AxoSimModelProfile", "kind": "class", "page": "api/model-family", "line": 25, "end_line": 35, "signature": "AxoSimModelProfile(profile_id: str, public_name: str, family: ModelFamily, model_dim: int, state_dim: int, num_layers: int, head_dim: int, mamba_update_normalization: Literal['none', 'fixed_rms'] = 'none') -> None", "parameters": [ { "name": "profile_id", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "public_name", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "family", "type": "ModelFamily", "default": null, "kind": "positional_or_keyword" }, { "name": "model_dim", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "state_dim", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "num_layers", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "head_dim", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "mamba_update_normalization", "type": "Literal['none', 'fixed_rms']", "default": "'none'", "kind": "positional_or_keyword" } ], "fields": [ { "name": "profile_id", "type": "str", "default": null }, { "name": "public_name", "type": "str", "default": null }, { "name": "family", "type": "ModelFamily", "default": null }, { "name": "model_dim", "type": "int", "default": null }, { "name": "state_dim", "type": "int", "default": null }, { "name": "num_layers", "type": "int", "default": null }, { "name": "head_dim", "type": "int", "default": null }, { "name": "mamba_update_normalization", "type": "Literal['none', 'fixed_rms']", "default": "'none'" } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "model_family", "name": "AXOSIM_MODEL_PROFILES", "kind": "constant", "page": "api/model-family", "line": 86, "end_line": 88, "signature": "AXOSIM_MODEL_PROFILES: Mapping[str, AxoSimModelProfile] = MappingProxyType({profile.profile_id: profile for profile in _PROFILES})", "parameters": [], "fields": [], "methods": [] }, { "module": "model_family", "name": "create_axosim_profile", "kind": "function", "page": "api/model-family", "line": 91, "end_line": 126, "signature": "create_axosim_profile(profile_id: str, *, base_config: AxoMambaConfig | None=None, use_pytorch_fallback: bool=False) -> nn.Module", "parameters": [ { "name": "profile_id", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "base_config", "type": "AxoMambaConfig | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "use_pytorch_fallback", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false } ], "fields": [], "methods": [] }, { "module": "axomamba", "name": "AxoMambaConfig", "kind": "class", "page": "api/mamba", "line": 20, "end_line": 21, "signature": "AxoMambaConfig(num_input: int = 1278, num_output: int = 2, num_branch: int = 45, num_synapse_per_branch: int = 100, input_to_synapse_routing: str | None = 'neuronio_routing', model_dim: int = 64, num_layers: int = 2, block_repeats: int = 1, state_dim: int = 16, conv_kernel: int = 4, expansion: int = 2, block_type: str = 'mamba1', head_dim: int = 64, num_groups: int = 1, chunk_size: int = 256, block_norm: bool = True, final_norm: bool = True, dropout: float = 0.0, residual_scale_init: float = 0.1, mamba_update_normalization: str = 'none', separate_heads: bool = False, soma_filter: bool = False, soma_filter_tau: float = 25.0, soma_filter_kernel: int = 129, learn_soma_filter_decay: bool = True, soma_multiplicative_gate: bool = False, soma_multiplicative_scale_init: float = 0.1, soma_peak_correction: bool = False, soma_peak_correction_scale_init: float = 0.1, soma_highpass_correction: bool = False, soma_highpass_correction_scale_init: float = 0.1, spike_voltage_coupling_scale: float = 0.0, morphology_ids: list[str] | None = None, morphology_embedding_scale: float | None = None, population_adapter_morphology_id: str | None = None, synapse_gain_scale: float | None = None, morphology_synapse_gain_scale: float | None = None, share_morphology_synapse_gain: bool = False, morphology_synapse_gain_rank: int = 0, morphology_synapse_feature_dim: int = 0, morphology_synapse_feature_hidden: int = 0, morphology_synapse_feature_scale: float | None = None, morphology_synapse_feature_storage_dtype: str = 'float32', local_tcn_scale: float | None = None, local_tcn_mode: str = 'dilated', local_tcn_position: str = 'parallel', share_local_tcn_pointwise: bool = False, local_tcn_pointwise_adapter_rank: int = 0, local_tcn_pointwise_rank: int = 0, local_tcn_output_rank: int = 0, mamba_projection_rank: int = 0, branch_gain: bool = False, branch_nonlinearity: str = 'none', branch_nonlinearity_scale_init: float = 0.1, branch_subunits: int = 0, branch_subunit_routing: str = 'contiguous', branch_subunit_scale: float = 1.0, branch_bilinear_rank: int = 0, branch_bilinear_scale: float | None = None, branch_trace_taus: list[float] | None = None, branch_trace_kernel: int = 65, hidden_trace_taus: list[float] | None = None, hidden_trace_kernel: int = 129, temporal_refine_kernel: int = 0, integrative_mixer: bool = False, integrative_mixer_short_kernel: int = 9, integrative_mixer_long_kernel: int = 65, integrative_mixer_position: str = 'pre', integrative_mixer_long_path_bias: float = 0.5) -> None", "parameters": [ { "name": "num_input", "type": "int", "default": "1278", "kind": "positional_or_keyword" }, { "name": "num_output", "type": "int", "default": "2", "kind": "positional_or_keyword" }, { "name": "num_branch", "type": "int", "default": "45", "kind": "positional_or_keyword" }, { "name": "num_synapse_per_branch", "type": "int", "default": "100", "kind": "positional_or_keyword" }, { "name": "input_to_synapse_routing", "type": "str | None", "default": "'neuronio_routing'", "kind": "positional_or_keyword" }, { "name": "model_dim", "type": "int", "default": "64", "kind": "positional_or_keyword" }, { "name": "num_layers", "type": "int", "default": "2", "kind": "positional_or_keyword" }, { "name": "block_repeats", "type": "int", "default": "1", "kind": "positional_or_keyword" }, { "name": "state_dim", "type": "int", "default": "16", "kind": "positional_or_keyword" }, { "name": "conv_kernel", "type": "int", "default": "4", "kind": "positional_or_keyword" }, { "name": "expansion", "type": "int", "default": "2", "kind": "positional_or_keyword" }, { "name": "block_type", "type": "str", "default": "'mamba1'", "kind": "positional_or_keyword" }, { "name": "head_dim", "type": "int", "default": "64", "kind": "positional_or_keyword" }, { "name": "num_groups", "type": "int", "default": "1", "kind": "positional_or_keyword" }, { "name": "chunk_size", "type": "int", "default": "256", "kind": "positional_or_keyword" }, { "name": "block_norm", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "final_norm", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "dropout", "type": "float", "default": "0.0", "kind": "positional_or_keyword" }, { "name": "residual_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "mamba_update_normalization", "type": "str", "default": "'none'", "kind": "positional_or_keyword" }, { "name": "separate_heads", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "soma_filter", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "soma_filter_tau", "type": "float", "default": "25.0", "kind": "positional_or_keyword" }, { "name": "soma_filter_kernel", "type": "int", "default": "129", "kind": "positional_or_keyword" }, { "name": "learn_soma_filter_decay", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "soma_multiplicative_gate", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "soma_multiplicative_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "soma_peak_correction", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "soma_peak_correction_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "soma_highpass_correction", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "soma_highpass_correction_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "spike_voltage_coupling_scale", "type": "float", "default": "0.0", "kind": "positional_or_keyword" }, { "name": "morphology_ids", "type": "list[str] | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "morphology_embedding_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "population_adapter_morphology_id", "type": "str | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "synapse_gain_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_gain_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "share_morphology_synapse_gain", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_gain_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_feature_dim", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_feature_hidden", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_feature_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_feature_storage_dtype", "type": "str", "default": "'float32'", "kind": "positional_or_keyword" }, { "name": "local_tcn_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "local_tcn_mode", "type": "str", "default": "'dilated'", "kind": "positional_or_keyword" }, { "name": "local_tcn_position", "type": "str", "default": "'parallel'", "kind": "positional_or_keyword" }, { "name": "share_local_tcn_pointwise", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "local_tcn_pointwise_adapter_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "local_tcn_pointwise_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "local_tcn_output_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "mamba_projection_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "branch_gain", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "branch_nonlinearity", "type": "str", "default": "'none'", "kind": "positional_or_keyword" }, { "name": "branch_nonlinearity_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "branch_subunits", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "branch_subunit_routing", "type": "str", "default": "'contiguous'", "kind": "positional_or_keyword" }, { "name": "branch_subunit_scale", "type": "float", "default": "1.0", "kind": "positional_or_keyword" }, { "name": "branch_bilinear_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "branch_bilinear_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "branch_trace_taus", "type": "list[float] | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "branch_trace_kernel", "type": "int", "default": "65", "kind": "positional_or_keyword" }, { "name": "hidden_trace_taus", "type": "list[float] | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "hidden_trace_kernel", "type": "int", "default": "129", "kind": "positional_or_keyword" }, { "name": "temporal_refine_kernel", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "integrative_mixer", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "integrative_mixer_short_kernel", "type": "int", "default": "9", "kind": "positional_or_keyword" }, { "name": "integrative_mixer_long_kernel", "type": "int", "default": "65", "kind": "positional_or_keyword" }, { "name": "integrative_mixer_position", "type": "str", "default": "'pre'", "kind": "positional_or_keyword" }, { "name": "integrative_mixer_long_path_bias", "type": "float", "default": "0.5", "kind": "positional_or_keyword" } ], "fields": [], "methods": [], "bases": [ "BranchOfficialMambaConfig" ], "inherited_fields": [ { "name": "num_input", "type": "int", "default": "1278" }, { "name": "num_output", "type": "int", "default": "2" }, { "name": "num_branch", "type": "int", "default": "45" }, { "name": "num_synapse_per_branch", "type": "int", "default": "100" }, { "name": "input_to_synapse_routing", "type": "str | None", "default": "'neuronio_routing'" }, { "name": "model_dim", "type": "int", "default": "64" }, { "name": "num_layers", "type": "int", "default": "2" }, { "name": "block_repeats", "type": "int", "default": "1" }, { "name": "state_dim", "type": "int", "default": "16" }, { "name": "conv_kernel", "type": "int", "default": "4" }, { "name": "expansion", "type": "int", "default": "2" }, { "name": "block_type", "type": "str", "default": "'mamba1'" }, { "name": "head_dim", "type": "int", "default": "64" }, { "name": "num_groups", "type": "int", "default": "1" }, { "name": "chunk_size", "type": "int", "default": "256" }, { "name": "block_norm", "type": "bool", "default": "True" }, { "name": "final_norm", "type": "bool", "default": "True" }, { "name": "dropout", "type": "float", "default": "0.0" }, { "name": "residual_scale_init", "type": "float", "default": "0.1" }, { "name": "mamba_update_normalization", "type": "str", "default": "'none'" }, { "name": "separate_heads", "type": "bool", "default": "False" }, { "name": "soma_filter", "type": "bool", "default": "False" }, { "name": "soma_filter_tau", "type": "float", "default": "25.0" }, { "name": "soma_filter_kernel", "type": "int", "default": "129" }, { "name": "learn_soma_filter_decay", "type": "bool", "default": "True" }, { "name": "soma_multiplicative_gate", "type": "bool", "default": "False" }, { "name": "soma_multiplicative_scale_init", "type": "float", "default": "0.1" }, { "name": "soma_peak_correction", "type": "bool", "default": "False" }, { "name": "soma_peak_correction_scale_init", "type": "float", "default": "0.1" }, { "name": "soma_highpass_correction", "type": "bool", "default": "False" }, { "name": "soma_highpass_correction_scale_init", "type": "float", "default": "0.1" }, { "name": "spike_voltage_coupling_scale", "type": "float", "default": "0.0" }, { "name": "morphology_ids", "type": "list[str] | None", "default": "None" }, { "name": "morphology_embedding_scale", "type": "float | None", "default": "None" }, { "name": "population_adapter_morphology_id", "type": "str | None", "default": "None" }, { "name": "synapse_gain_scale", "type": "float | None", "default": "None" }, { "name": "morphology_synapse_gain_scale", "type": "float | None", "default": "None" }, { "name": "share_morphology_synapse_gain", "type": "bool", "default": "False" }, { "name": "morphology_synapse_gain_rank", "type": "int", "default": "0" }, { "name": "morphology_synapse_feature_dim", "type": "int", "default": "0" }, { "name": "morphology_synapse_feature_hidden", "type": "int", "default": "0" }, { "name": "morphology_synapse_feature_scale", "type": "float | None", "default": "None" }, { "name": "morphology_synapse_feature_storage_dtype", "type": "str", "default": "'float32'" }, { "name": "local_tcn_scale", "type": "float | None", "default": "None" }, { "name": "local_tcn_mode", "type": "str", "default": "'dilated'" }, { "name": "local_tcn_position", "type": "str", "default": "'parallel'" }, { "name": "share_local_tcn_pointwise", "type": "bool", "default": "False" }, { "name": "local_tcn_pointwise_adapter_rank", "type": "int", "default": "0" }, { "name": "local_tcn_pointwise_rank", "type": "int", "default": "0" }, { "name": "local_tcn_output_rank", "type": "int", "default": "0" }, { "name": "mamba_projection_rank", "type": "int", "default": "0" }, { "name": "branch_gain", "type": "bool", "default": "False" }, { "name": "branch_nonlinearity", "type": "str", "default": "'none'" }, { "name": "branch_nonlinearity_scale_init", "type": "float", "default": "0.1" }, { "name": "branch_subunits", "type": "int", "default": "0" }, { "name": "branch_subunit_routing", "type": "str", "default": "'contiguous'" }, { "name": "branch_subunit_scale", "type": "float", "default": "1.0" }, { "name": "branch_bilinear_rank", "type": "int", "default": "0" }, { "name": "branch_bilinear_scale", "type": "float | None", "default": "None" }, { "name": "branch_trace_taus", "type": "list[float] | None", "default": "None" }, { "name": "branch_trace_kernel", "type": "int", "default": "65" }, { "name": "hidden_trace_taus", "type": "list[float] | None", "default": "None" }, { "name": "hidden_trace_kernel", "type": "int", "default": "129" }, { "name": "temporal_refine_kernel", "type": "int", "default": "0" }, { "name": "integrative_mixer", "type": "bool", "default": "False" }, { "name": "integrative_mixer_short_kernel", "type": "int", "default": "9" }, { "name": "integrative_mixer_long_kernel", "type": "int", "default": "65" }, { "name": "integrative_mixer_position", "type": "str", "default": "'pre'" }, { "name": "integrative_mixer_long_path_bias", "type": "float", "default": "0.5" } ] }, { "module": "axomamba", "name": "default_axomamba_config", "kind": "function", "page": "api/mamba", "line": 24, "end_line": 83, "signature": "default_axomamba_config(**overrides: Any) -> AxoMambaConfig", "parameters": [ { "name": "overrides", "type": "Any", "default": null, "kind": "variadic_keyword" } ], "fields": [], "methods": [] }, { "module": "axomamba", "name": "structured_compact_axomamba_config", "kind": "function", "page": "api/mamba", "line": 86, "end_line": 96, "signature": "structured_compact_axomamba_config(**overrides: Any) -> AxoMambaConfig", "parameters": [ { "name": "overrides", "type": "Any", "default": null, "kind": "variadic_keyword" } ], "fields": [], "methods": [] }, { "module": "axomamba", "name": "regression_axomamba_config", "kind": "function", "page": "api/mamba", "line": 99, "end_line": 108, "signature": "regression_axomamba_config(**overrides: Any) -> AxoMambaConfig", "parameters": [ { "name": "overrides", "type": "Any", "default": null, "kind": "variadic_keyword" } ], "fields": [], "methods": [] }, { "module": "axomamba", "name": "population_axomamba_config", "kind": "function", "page": "api/mamba", "line": 111, "end_line": 125, "signature": "population_axomamba_config(**overrides: Any) -> AxoMambaConfig", "parameters": [ { "name": "overrides", "type": "Any", "default": null, "kind": "variadic_keyword" } ], "fields": [], "methods": [] }, { "module": "axomamba", "name": "spike_axomamba_config", "kind": "function", "page": "api/mamba", "line": 128, "end_line": 139, "signature": "spike_axomamba_config(**overrides: Any) -> AxoMambaConfig", "parameters": [ { "name": "overrides", "type": "Any", "default": null, "kind": "variadic_keyword" } ], "fields": [], "methods": [] }, { "module": "axomamba", "name": "load_axomamba_config", "kind": "function", "page": "api/mamba", "line": 142, "end_line": 148, "signature": "load_axomamba_config(path: str | Path | None=None, **overrides: Any) -> AxoMambaConfig", "parameters": [ { "name": "path", "type": "str | Path | None", "default": "None", "kind": "positional_or_keyword", "resolved_default": null }, { "name": "overrides", "type": "Any", "default": null, "kind": "variadic_keyword" } ], "fields": [], "methods": [] }, { "module": "axomamba", "name": "coerce_axomamba_config", "kind": "function", "page": "api/mamba", "line": 151, "end_line": 158, "signature": "coerce_axomamba_config(config: AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None) -> AxoMambaConfig", "parameters": [ { "name": "config", "type": "AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "axomamba", "name": "AxoMamba", "kind": "class", "page": "api/mamba", "line": 161, "end_line": 170, "signature": "AxoMamba(self, config: AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None=None, *, mamba_classes: tuple[type, type | None] | None=None) -> None", "parameters": [ { "name": "config", "type": "AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None", "default": "None", "kind": "positional_or_keyword", "resolved_default": null }, { "name": "mamba_classes", "type": "tuple[type, type | None] | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, config: AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None=None, *, mamba_classes: tuple[type, type | None] | None=None) -> None", "parameters": [ { "name": "config", "type": "AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None", "default": "None", "kind": "positional_or_keyword", "resolved_default": null }, { "name": "mamba_classes", "type": "tuple[type, type | None] | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 164, "end_line": 170, "return_type": "None", "decorators": [] } ], "bases": [ "BranchOfficialMamba" ] }, { "module": "axomamba", "name": "AxoPyTorchMamba", "kind": "class", "page": "api/mamba", "line": 173, "end_line": 182, "signature": "AxoPyTorchMamba(self, config: AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None=None) -> None", "parameters": [ { "name": "config", "type": "AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None", "default": "None", "kind": "positional_or_keyword", "resolved_default": null } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, config: AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None=None) -> None", "parameters": [ { "name": "config", "type": "AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None", "default": "None", "kind": "positional_or_keyword", "resolved_default": null } ], "line": 176, "end_line": 182, "return_type": "None", "decorators": [] } ], "bases": [ "AxoMamba" ] }, { "module": "axomamba", "name": "create_axomamba", "kind": "function", "page": "api/mamba", "line": 185, "end_line": 192, "signature": "create_axomamba(config: AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None=None, *, use_pytorch_fallback: bool=False) -> AxoMamba", "parameters": [ { "name": "config", "type": "AxoMambaConfig | BranchOfficialMambaConfig | Mapping[str, Any] | None", "default": "None", "kind": "positional_or_keyword", "resolved_default": null }, { "name": "use_pytorch_fallback", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false } ], "fields": [], "methods": [] }, { "module": "mamba_official", "name": "BranchOfficialMambaConfig", "kind": "class", "page": "api/mamba", "line": 29, "end_line": 98, "signature": "BranchOfficialMambaConfig(num_input: int = 1278, num_output: int = 2, num_branch: int = 45, num_synapse_per_branch: int = 100, input_to_synapse_routing: str | None = 'neuronio_routing', model_dim: int = 64, num_layers: int = 2, block_repeats: int = 1, state_dim: int = 16, conv_kernel: int = 4, expansion: int = 2, block_type: str = 'mamba1', head_dim: int = 64, num_groups: int = 1, chunk_size: int = 256, block_norm: bool = True, final_norm: bool = True, dropout: float = 0.0, residual_scale_init: float = 0.1, mamba_update_normalization: str = 'none', separate_heads: bool = False, soma_filter: bool = False, soma_filter_tau: float = 25.0, soma_filter_kernel: int = 129, learn_soma_filter_decay: bool = True, soma_multiplicative_gate: bool = False, soma_multiplicative_scale_init: float = 0.1, soma_peak_correction: bool = False, soma_peak_correction_scale_init: float = 0.1, soma_highpass_correction: bool = False, soma_highpass_correction_scale_init: float = 0.1, spike_voltage_coupling_scale: float = 0.0, morphology_ids: list[str] | None = None, morphology_embedding_scale: float | None = None, population_adapter_morphology_id: str | None = None, synapse_gain_scale: float | None = None, morphology_synapse_gain_scale: float | None = None, share_morphology_synapse_gain: bool = False, morphology_synapse_gain_rank: int = 0, morphology_synapse_feature_dim: int = 0, morphology_synapse_feature_hidden: int = 0, morphology_synapse_feature_scale: float | None = None, morphology_synapse_feature_storage_dtype: str = 'float32', local_tcn_scale: float | None = None, local_tcn_mode: str = 'dilated', local_tcn_position: str = 'parallel', share_local_tcn_pointwise: bool = False, local_tcn_pointwise_adapter_rank: int = 0, local_tcn_pointwise_rank: int = 0, local_tcn_output_rank: int = 0, mamba_projection_rank: int = 0, branch_gain: bool = False, branch_nonlinearity: str = 'none', branch_nonlinearity_scale_init: float = 0.1, branch_subunits: int = 0, branch_subunit_routing: str = 'contiguous', branch_subunit_scale: float = 1.0, branch_bilinear_rank: int = 0, branch_bilinear_scale: float | None = None, branch_trace_taus: list[float] | None = None, branch_trace_kernel: int = 65, hidden_trace_taus: list[float] | None = None, hidden_trace_kernel: int = 129, temporal_refine_kernel: int = 0, integrative_mixer: bool = False, integrative_mixer_short_kernel: int = 9, integrative_mixer_long_kernel: int = 65, integrative_mixer_position: str = 'pre', integrative_mixer_long_path_bias: float = 0.5) -> None", "parameters": [ { "name": "num_input", "type": "int", "default": "1278", "kind": "positional_or_keyword" }, { "name": "num_output", "type": "int", "default": "2", "kind": "positional_or_keyword" }, { "name": "num_branch", "type": "int", "default": "45", "kind": "positional_or_keyword" }, { "name": "num_synapse_per_branch", "type": "int", "default": "100", "kind": "positional_or_keyword" }, { "name": "input_to_synapse_routing", "type": "str | None", "default": "'neuronio_routing'", "kind": "positional_or_keyword" }, { "name": "model_dim", "type": "int", "default": "64", "kind": "positional_or_keyword" }, { "name": "num_layers", "type": "int", "default": "2", "kind": "positional_or_keyword" }, { "name": "block_repeats", "type": "int", "default": "1", "kind": "positional_or_keyword" }, { "name": "state_dim", "type": "int", "default": "16", "kind": "positional_or_keyword" }, { "name": "conv_kernel", "type": "int", "default": "4", "kind": "positional_or_keyword" }, { "name": "expansion", "type": "int", "default": "2", "kind": "positional_or_keyword" }, { "name": "block_type", "type": "str", "default": "'mamba1'", "kind": "positional_or_keyword" }, { "name": "head_dim", "type": "int", "default": "64", "kind": "positional_or_keyword" }, { "name": "num_groups", "type": "int", "default": "1", "kind": "positional_or_keyword" }, { "name": "chunk_size", "type": "int", "default": "256", "kind": "positional_or_keyword" }, { "name": "block_norm", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "final_norm", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "dropout", "type": "float", "default": "0.0", "kind": "positional_or_keyword" }, { "name": "residual_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "mamba_update_normalization", "type": "str", "default": "'none'", "kind": "positional_or_keyword" }, { "name": "separate_heads", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "soma_filter", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "soma_filter_tau", "type": "float", "default": "25.0", "kind": "positional_or_keyword" }, { "name": "soma_filter_kernel", "type": "int", "default": "129", "kind": "positional_or_keyword" }, { "name": "learn_soma_filter_decay", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "soma_multiplicative_gate", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "soma_multiplicative_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "soma_peak_correction", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "soma_peak_correction_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "soma_highpass_correction", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "soma_highpass_correction_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "spike_voltage_coupling_scale", "type": "float", "default": "0.0", "kind": "positional_or_keyword" }, { "name": "morphology_ids", "type": "list[str] | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "morphology_embedding_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "population_adapter_morphology_id", "type": "str | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "synapse_gain_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_gain_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "share_morphology_synapse_gain", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_gain_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_feature_dim", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_feature_hidden", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_feature_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "morphology_synapse_feature_storage_dtype", "type": "str", "default": "'float32'", "kind": "positional_or_keyword" }, { "name": "local_tcn_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "local_tcn_mode", "type": "str", "default": "'dilated'", "kind": "positional_or_keyword" }, { "name": "local_tcn_position", "type": "str", "default": "'parallel'", "kind": "positional_or_keyword" }, { "name": "share_local_tcn_pointwise", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "local_tcn_pointwise_adapter_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "local_tcn_pointwise_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "local_tcn_output_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "mamba_projection_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "branch_gain", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "branch_nonlinearity", "type": "str", "default": "'none'", "kind": "positional_or_keyword" }, { "name": "branch_nonlinearity_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "branch_subunits", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "branch_subunit_routing", "type": "str", "default": "'contiguous'", "kind": "positional_or_keyword" }, { "name": "branch_subunit_scale", "type": "float", "default": "1.0", "kind": "positional_or_keyword" }, { "name": "branch_bilinear_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "branch_bilinear_scale", "type": "float | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "branch_trace_taus", "type": "list[float] | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "branch_trace_kernel", "type": "int", "default": "65", "kind": "positional_or_keyword" }, { "name": "hidden_trace_taus", "type": "list[float] | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "hidden_trace_kernel", "type": "int", "default": "129", "kind": "positional_or_keyword" }, { "name": "temporal_refine_kernel", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "integrative_mixer", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "integrative_mixer_short_kernel", "type": "int", "default": "9", "kind": "positional_or_keyword" }, { "name": "integrative_mixer_long_kernel", "type": "int", "default": "65", "kind": "positional_or_keyword" }, { "name": "integrative_mixer_position", "type": "str", "default": "'pre'", "kind": "positional_or_keyword" }, { "name": "integrative_mixer_long_path_bias", "type": "float", "default": "0.5", "kind": "positional_or_keyword" } ], "fields": [ { "name": "num_input", "type": "int", "default": "1278" }, { "name": "num_output", "type": "int", "default": "2" }, { "name": "num_branch", "type": "int", "default": "45" }, { "name": "num_synapse_per_branch", "type": "int", "default": "100" }, { "name": "input_to_synapse_routing", "type": "str | None", "default": "'neuronio_routing'" }, { "name": "model_dim", "type": "int", "default": "64" }, { "name": "num_layers", "type": "int", "default": "2" }, { "name": "block_repeats", "type": "int", "default": "1" }, { "name": "state_dim", "type": "int", "default": "16" }, { "name": "conv_kernel", "type": "int", "default": "4" }, { "name": "expansion", "type": "int", "default": "2" }, { "name": "block_type", "type": "str", "default": "'mamba1'" }, { "name": "head_dim", "type": "int", "default": "64" }, { "name": "num_groups", "type": "int", "default": "1" }, { "name": "chunk_size", "type": "int", "default": "256" }, { "name": "block_norm", "type": "bool", "default": "True" }, { "name": "final_norm", "type": "bool", "default": "True" }, { "name": "dropout", "type": "float", "default": "0.0" }, { "name": "residual_scale_init", "type": "float", "default": "0.1" }, { "name": "mamba_update_normalization", "type": "str", "default": "'none'" }, { "name": "separate_heads", "type": "bool", "default": "False" }, { "name": "soma_filter", "type": "bool", "default": "False" }, { "name": "soma_filter_tau", "type": "float", "default": "25.0" }, { "name": "soma_filter_kernel", "type": "int", "default": "129" }, { "name": "learn_soma_filter_decay", "type": "bool", "default": "True" }, { "name": "soma_multiplicative_gate", "type": "bool", "default": "False" }, { "name": "soma_multiplicative_scale_init", "type": "float", "default": "0.1" }, { "name": "soma_peak_correction", "type": "bool", "default": "False" }, { "name": "soma_peak_correction_scale_init", "type": "float", "default": "0.1" }, { "name": "soma_highpass_correction", "type": "bool", "default": "False" }, { "name": "soma_highpass_correction_scale_init", "type": "float", "default": "0.1" }, { "name": "spike_voltage_coupling_scale", "type": "float", "default": "0.0" }, { "name": "morphology_ids", "type": "list[str] | None", "default": "None" }, { "name": "morphology_embedding_scale", "type": "float | None", "default": "None" }, { "name": "population_adapter_morphology_id", "type": "str | None", "default": "None" }, { "name": "synapse_gain_scale", "type": "float | None", "default": "None" }, { "name": "morphology_synapse_gain_scale", "type": "float | None", "default": "None" }, { "name": "share_morphology_synapse_gain", "type": "bool", "default": "False" }, { "name": "morphology_synapse_gain_rank", "type": "int", "default": "0" }, { "name": "morphology_synapse_feature_dim", "type": "int", "default": "0" }, { "name": "morphology_synapse_feature_hidden", "type": "int", "default": "0" }, { "name": "morphology_synapse_feature_scale", "type": "float | None", "default": "None" }, { "name": "morphology_synapse_feature_storage_dtype", "type": "str", "default": "'float32'" }, { "name": "local_tcn_scale", "type": "float | None", "default": "None" }, { "name": "local_tcn_mode", "type": "str", "default": "'dilated'" }, { "name": "local_tcn_position", "type": "str", "default": "'parallel'" }, { "name": "share_local_tcn_pointwise", "type": "bool", "default": "False" }, { "name": "local_tcn_pointwise_adapter_rank", "type": "int", "default": "0" }, { "name": "local_tcn_pointwise_rank", "type": "int", "default": "0" }, { "name": "local_tcn_output_rank", "type": "int", "default": "0" }, { "name": "mamba_projection_rank", "type": "int", "default": "0" }, { "name": "branch_gain", "type": "bool", "default": "False" }, { "name": "branch_nonlinearity", "type": "str", "default": "'none'" }, { "name": "branch_nonlinearity_scale_init", "type": "float", "default": "0.1" }, { "name": "branch_subunits", "type": "int", "default": "0" }, { "name": "branch_subunit_routing", "type": "str", "default": "'contiguous'" }, { "name": "branch_subunit_scale", "type": "float", "default": "1.0" }, { "name": "branch_bilinear_rank", "type": "int", "default": "0" }, { "name": "branch_bilinear_scale", "type": "float | None", "default": "None" }, { "name": "branch_trace_taus", "type": "list[float] | None", "default": "None" }, { "name": "branch_trace_kernel", "type": "int", "default": "65" }, { "name": "hidden_trace_taus", "type": "list[float] | None", "default": "None" }, { "name": "hidden_trace_kernel", "type": "int", "default": "129" }, { "name": "temporal_refine_kernel", "type": "int", "default": "0" }, { "name": "integrative_mixer", "type": "bool", "default": "False" }, { "name": "integrative_mixer_short_kernel", "type": "int", "default": "9" }, { "name": "integrative_mixer_long_kernel", "type": "int", "default": "65" }, { "name": "integrative_mixer_position", "type": "str", "default": "'pre'" }, { "name": "integrative_mixer_long_path_bias", "type": "float", "default": "0.5" } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "mamba_official", "name": "BranchMambaStreamingState", "kind": "class", "page": "api/mamba", "line": 102, "end_line": 107, "signature": "BranchMambaStreamingState(mamba_states: list[tuple[torch.Tensor, torch.Tensor]], local_tcn_histories: list[torch.Tensor], morphology_feature_gains: torch.Tensor | None = None) -> None", "parameters": [ { "name": "mamba_states", "type": "list[tuple[torch.Tensor, torch.Tensor]]", "default": null, "kind": "positional_or_keyword" }, { "name": "local_tcn_histories", "type": "list[torch.Tensor]", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_feature_gains", "type": "torch.Tensor | None", "default": "None", "kind": "positional_or_keyword" } ], "fields": [ { "name": "mamba_states", "type": "list[tuple[torch.Tensor, torch.Tensor]]", "default": null }, { "name": "local_tcn_histories", "type": "list[torch.Tensor]", "default": null }, { "name": "morphology_feature_gains", "type": "torch.Tensor | None", "default": "None" } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "mamba_official", "name": "BranchOfficialMamba", "kind": "class", "page": "api/mamba", "line": 110, "end_line": 1808, "signature": "BranchOfficialMamba(self, config: BranchOfficialMambaConfig, *, mamba_classes: tuple[type[nn.Module], type[nn.Module] | None] | None=None) -> None", "parameters": [ { "name": "config", "type": "BranchOfficialMambaConfig", "default": null, "kind": "positional_or_keyword" }, { "name": "mamba_classes", "type": "tuple[type[nn.Module], type[nn.Module] | None] | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, config: BranchOfficialMambaConfig, *, mamba_classes: tuple[type[nn.Module], type[nn.Module] | None] | None=None) -> None", "parameters": [ { "name": "config", "type": "BranchOfficialMambaConfig", "default": null, "kind": "positional_or_keyword" }, { "name": "mamba_classes", "type": "tuple[type[nn.Module], type[nn.Module] | None] | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 113, "end_line": 588, "return_type": "None", "decorators": [] }, { "name": "forward", "kind": "method", "signature": "forward(self, x: torch.Tensor, *, morphology_indices: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "x", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 678, "end_line": 689, "return_type": "torch.Tensor", "decorators": [] }, { "name": "allocate_streaming_state", "kind": "method", "signature": "allocate_streaming_state(self, batch_size: int, *, device: torch.device | str | None=None, dtype: torch.dtype | None=None) -> BranchMambaStreamingState", "parameters": [ { "name": "batch_size", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "device", "type": "torch.device | str | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "dtype", "type": "torch.dtype | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 954, "end_line": 994, "return_type": "BranchMambaStreamingState", "decorators": [] }, { "name": "streaming_step", "kind": "method", "signature": "streaming_step(self, x: torch.Tensor, state: BranchMambaStreamingState, *, morphology_indices: torch.Tensor | None=None, chunk_size: int | None=None) -> torch.Tensor", "parameters": [ { "name": "x", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "state", "type": "BranchMambaStreamingState", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "chunk_size", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 996, "end_line": 1087, "return_type": "torch.Tensor", "decorators": [] }, { "name": "streaming_step_events", "kind": "method", "signature": "streaming_step_events(self, event_indices: torch.Tensor, event_values: torch.Tensor, state: BranchMambaStreamingState, *, morphology_indices: torch.Tensor | None=None, chunk_size: int | None=None, output_buffer: torch.Tensor | None=None, retain_base_soma_prediction: bool=True) -> torch.Tensor", "parameters": [ { "name": "event_indices", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "event_values", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "state", "type": "BranchMambaStreamingState", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "chunk_size", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "output_buffer", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "retain_base_soma_prediction", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "line": 1310, "end_line": 1451, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [ "nn.Module" ] }, { "module": "mamba_official", "name": "BranchPyTorchMamba", "kind": "class", "page": "api/mamba", "line": 1810, "end_line": 1818, "signature": "BranchPyTorchMamba(self, config: BranchOfficialMambaConfig) -> None", "parameters": [ { "name": "config", "type": "BranchOfficialMambaConfig", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, config: BranchOfficialMambaConfig) -> None", "parameters": [ { "name": "config", "type": "BranchOfficialMambaConfig", "default": null, "kind": "positional_or_keyword" } ], "line": 1813, "end_line": 1818, "return_type": "None", "decorators": [] } ], "bases": [ "BranchOfficialMamba" ] }, { "module": "temporal_core", "name": "AxoTemporalStreamingState", "kind": "class", "page": "api/temporal-core", "line": 16, "end_line": 24, "signature": "AxoTemporalStreamingState(core_state: torch.Tensor | tuple[torch.Tensor, torch.Tensor], local_tcn_histories: list[torch.Tensor], morphology_feature_gains: torch.Tensor | None, patch_sum: torch.Tensor | None = None, patch_correction: torch.Tensor | None = None, patch_position: int = 0) -> None", "parameters": [ { "name": "core_state", "type": "torch.Tensor | tuple[torch.Tensor, torch.Tensor]", "default": null, "kind": "positional_or_keyword" }, { "name": "local_tcn_histories", "type": "list[torch.Tensor]", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_feature_gains", "type": "torch.Tensor | None", "default": null, "kind": "positional_or_keyword" }, { "name": "patch_sum", "type": "torch.Tensor | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "patch_correction", "type": "torch.Tensor | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "patch_position", "type": "int", "default": "0", "kind": "positional_or_keyword" } ], "fields": [ { "name": "core_state", "type": "torch.Tensor | tuple[torch.Tensor, torch.Tensor]", "default": null }, { "name": "local_tcn_histories", "type": "list[torch.Tensor]", "default": null }, { "name": "morphology_feature_gains", "type": "torch.Tensor | None", "default": null }, { "name": "patch_sum", "type": "torch.Tensor | None", "default": "None" }, { "name": "patch_correction", "type": "torch.Tensor | None", "default": "None" }, { "name": "patch_position", "type": "int", "default": "0" } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "temporal_core", "name": "AxoTemporalCoreConfig", "kind": "class", "page": "api/temporal-core", "line": 28, "end_line": 85, "signature": "AxoTemporalCoreConfig(kind: TemporalCoreKind, gru_hidden_units: int | None = None, branch_memory_units: int = 30, branch_hidden_units: int = 64, branch_synapse_decay: float = 0.85, branch_memory_decay: float = 0.9, residual_scale_init: float = 0.1, keep_local_tcn: bool = False, patch_size: int = 1, behavior_adapter_morphology_id: str | None = None, behavior_adapter_rank: int = 0, behavior_adapter_branch_token_offset: bool = False) -> None", "parameters": [ { "name": "kind", "type": "TemporalCoreKind", "default": null, "kind": "positional_or_keyword" }, { "name": "gru_hidden_units", "type": "int | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "branch_memory_units", "type": "int", "default": "30", "kind": "positional_or_keyword" }, { "name": "branch_hidden_units", "type": "int", "default": "64", "kind": "positional_or_keyword" }, { "name": "branch_synapse_decay", "type": "float", "default": "0.85", "kind": "positional_or_keyword" }, { "name": "branch_memory_decay", "type": "float", "default": "0.9", "kind": "positional_or_keyword" }, { "name": "residual_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "keep_local_tcn", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "patch_size", "type": "int", "default": "1", "kind": "positional_or_keyword" }, { "name": "behavior_adapter_morphology_id", "type": "str | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "behavior_adapter_rank", "type": "int", "default": "0", "kind": "positional_or_keyword" }, { "name": "behavior_adapter_branch_token_offset", "type": "bool", "default": "False", "kind": "positional_or_keyword" } ], "fields": [ { "name": "kind", "type": "TemporalCoreKind", "default": null }, { "name": "gru_hidden_units", "type": "int | None", "default": "None" }, { "name": "branch_memory_units", "type": "int", "default": "30" }, { "name": "branch_hidden_units", "type": "int", "default": "64" }, { "name": "branch_synapse_decay", "type": "float", "default": "0.85" }, { "name": "branch_memory_decay", "type": "float", "default": "0.9" }, { "name": "residual_scale_init", "type": "float", "default": "0.1" }, { "name": "keep_local_tcn", "type": "bool", "default": "False" }, { "name": "patch_size", "type": "int", "default": "1" }, { "name": "behavior_adapter_morphology_id", "type": "str | None", "default": "None" }, { "name": "behavior_adapter_rank", "type": "int", "default": "0" }, { "name": "behavior_adapter_branch_token_offset", "type": "bool", "default": "False" } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "temporal_core", "name": "NeuronBehaviorAdapter", "kind": "class", "page": "api/temporal-core", "line": 88, "end_line": 205, "signature": "NeuronBehaviorAdapter(self, *, branches: int, width: int, outputs: int, morphology_index: int, rank: int=0, branch_token_offset: bool=False) -> None", "parameters": [ { "name": "branches", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "width", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "outputs", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "morphology_index", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "rank", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "branch_token_offset", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, *, branches: int, width: int, outputs: int, morphology_index: int, rank: int=0, branch_token_offset: bool=False) -> None", "parameters": [ { "name": "branches", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "width", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "outputs", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "morphology_index", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "rank", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "branch_token_offset", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false } ], "line": 91, "end_line": 159, "return_type": "None", "decorators": [] }, { "name": "parameter_counts", "kind": "property", "signature": "NeuronBehaviorAdapter.parameter_counts: dict[str, int]", "parameters": [], "line": 162, "end_line": 201, "return_type": "dict[str, int]", "decorators": [ "property" ] }, { "name": "parameter_count", "kind": "property", "signature": "NeuronBehaviorAdapter.parameter_count: int", "parameters": [], "line": 204, "end_line": 205, "return_type": "int", "decorators": [ "property" ] } ], "bases": [ "nn.Module" ] }, { "module": "temporal_core", "name": "ResidualGRUTemporalCore", "kind": "class", "page": "api/temporal-core", "line": 208, "end_line": 328, "signature": "ResidualGRUTemporalCore(self, width: int, *, hidden_units: int | None=None, residual_scale_init: float) -> None", "parameters": [ { "name": "width", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "hidden_units", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "residual_scale_init", "type": "float", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, width: int, *, hidden_units: int | None=None, residual_scale_init: float) -> None", "parameters": [ { "name": "width", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "hidden_units", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "residual_scale_init", "type": "float", "default": null, "kind": "keyword_only" } ], "line": 211, "end_line": 238, "return_type": "None", "decorators": [] }, { "name": "forward", "kind": "method", "signature": "forward(self, hidden: torch.Tensor, *, weight_ih_delta: torch.Tensor | None=None, weight_hh_delta: torch.Tensor | None=None, adapter_mask: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "hidden", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "weight_ih_delta", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "weight_hh_delta", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "adapter_mask", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 240, "end_line": 253, "return_type": "torch.Tensor", "decorators": [] }, { "name": "temporal_correction", "kind": "method", "signature": "temporal_correction(self, hidden: torch.Tensor, *, weight_ih_delta: torch.Tensor | None=None, weight_hh_delta: torch.Tensor | None=None, adapter_mask: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "hidden", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "weight_ih_delta", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "weight_hh_delta", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "adapter_mask", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 255, "end_line": 286, "return_type": "torch.Tensor", "decorators": [] }, { "name": "recurrent_state_elements", "kind": "property", "signature": "ResidualGRUTemporalCore.recurrent_state_elements: int", "parameters": [], "line": 327, "end_line": 328, "return_type": "int", "decorators": [ "property" ] } ], "bases": [ "nn.Module" ] }, { "module": "temporal_core", "name": "BranchELMTemporalCore", "kind": "class", "page": "api/temporal-core", "line": 331, "end_line": 393, "signature": "BranchELMTemporalCore(self, width: int, *, memory_units: int, hidden_units: int, synapse_decay: float, memory_decay: float, residual_scale_init: float) -> None", "parameters": [ { "name": "width", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "memory_units", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "hidden_units", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "synapse_decay", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "memory_decay", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "residual_scale_init", "type": "float", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, width: int, *, memory_units: int, hidden_units: int, synapse_decay: float, memory_decay: float, residual_scale_init: float) -> None", "parameters": [ { "name": "width", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "memory_units", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "hidden_units", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "synapse_decay", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "memory_decay", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "residual_scale_init", "type": "float", "default": null, "kind": "keyword_only" } ], "line": 334, "end_line": 359, "return_type": "None", "decorators": [] }, { "name": "forward", "kind": "method", "signature": "forward(self, hidden: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "hidden", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 361, "end_line": 362, "return_type": "torch.Tensor", "decorators": [] }, { "name": "temporal_correction", "kind": "method", "signature": "temporal_correction(self, hidden: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "hidden", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 364, "end_line": 389, "return_type": "torch.Tensor", "decorators": [] }, { "name": "recurrent_state_elements", "kind": "property", "signature": "BranchELMTemporalCore.recurrent_state_elements: int", "parameters": [], "line": 392, "end_line": 393, "return_type": "int", "decorators": [ "property" ] } ], "bases": [ "nn.Module" ] }, { "module": "temporal_core", "name": "CausalPatchedTemporalCore", "kind": "class", "page": "api/temporal-core", "line": 396, "end_line": 455, "signature": "CausalPatchedTemporalCore(self, core: nn.Module, *, width: int, patch_size: int) -> None", "parameters": [ { "name": "core", "type": "nn.Module", "default": null, "kind": "positional_or_keyword" }, { "name": "width", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "patch_size", "type": "int", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, core: nn.Module, *, width: int, patch_size: int) -> None", "parameters": [ { "name": "core", "type": "nn.Module", "default": null, "kind": "positional_or_keyword" }, { "name": "width", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "patch_size", "type": "int", "default": null, "kind": "keyword_only" } ], "line": 399, "end_line": 411, "return_type": "None", "decorators": [] }, { "name": "forward", "kind": "method", "signature": "forward(self, hidden: torch.Tensor, **temporal_kwargs) -> torch.Tensor", "parameters": [ { "name": "hidden", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "temporal_kwargs", "type": null, "default": null, "kind": "variadic_keyword" } ], "line": 413, "end_line": 448, "return_type": "torch.Tensor", "decorators": [] }, { "name": "recurrent_state_elements", "kind": "property", "signature": "CausalPatchedTemporalCore.recurrent_state_elements: int", "parameters": [], "line": 451, "end_line": 455, "return_type": "int", "decorators": [ "property" ] } ], "bases": [ "nn.Module" ] }, { "module": "temporal_core", "name": "AxoTemporalModel", "kind": "class", "page": "api/temporal-core", "line": 458, "end_line": 1163, "signature": "AxoTemporalModel(self, source: BranchOfficialMamba, temporal_config: AxoTemporalCoreConfig) -> None", "parameters": [ { "name": "source", "type": "BranchOfficialMamba", "default": null, "kind": "positional_or_keyword" }, { "name": "temporal_config", "type": "AxoTemporalCoreConfig", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, source: BranchOfficialMamba, temporal_config: AxoTemporalCoreConfig) -> None", "parameters": [ { "name": "source", "type": "BranchOfficialMamba", "default": null, "kind": "positional_or_keyword" }, { "name": "temporal_config", "type": "AxoTemporalCoreConfig", "default": null, "kind": "positional_or_keyword" } ], "line": 461, "end_line": 522, "return_type": "None", "decorators": [] }, { "name": "config", "kind": "property", "signature": "AxoTemporalModel.config", "parameters": [], "line": 525, "end_line": 526, "return_type": null, "decorators": [ "property" ] }, { "name": "num_input", "kind": "property", "signature": "AxoTemporalModel.num_input: int", "parameters": [], "line": 529, "end_line": 530, "return_type": "int", "decorators": [ "property" ] }, { "name": "num_output", "kind": "property", "signature": "AxoTemporalModel.num_output: int", "parameters": [], "line": 533, "end_line": 534, "return_type": "int", "decorators": [ "property" ] }, { "name": "num_branch", "kind": "property", "signature": "AxoTemporalModel.num_branch: int", "parameters": [], "line": 537, "end_line": 538, "return_type": "int", "decorators": [ "property" ] }, { "name": "base_soma_prediction", "kind": "property", "signature": "AxoTemporalModel.base_soma_prediction: torch.Tensor | None", "parameters": [], "line": 541, "end_line": 542, "return_type": "torch.Tensor | None", "decorators": [ "property" ] }, { "name": "forward", "kind": "method", "signature": "forward(self, x: torch.Tensor, *, morphology_indices: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "x", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 544, "end_line": 652, "return_type": "torch.Tensor", "decorators": [] }, { "name": "allocate_streaming_state", "kind": "method", "signature": "allocate_streaming_state(self, batch_size: int, *, device: torch.device | str | None=None, dtype: torch.dtype | None=None) -> AxoTemporalStreamingState", "parameters": [ { "name": "batch_size", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "device", "type": "torch.device | str | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "dtype", "type": "torch.dtype | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 654, "end_line": 729, "return_type": "AxoTemporalStreamingState", "decorators": [] }, { "name": "streaming_step", "kind": "method", "signature": "streaming_step(self, x: torch.Tensor, state: AxoTemporalStreamingState, *, morphology_indices: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "x", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "state", "type": "AxoTemporalStreamingState", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 731, "end_line": 769, "return_type": "torch.Tensor", "decorators": [] }, { "name": "streaming_step_events", "kind": "method", "signature": "streaming_step_events(self, event_indices: torch.Tensor, event_values: torch.Tensor, state: AxoTemporalStreamingState, *, morphology_indices: torch.Tensor | None=None, chunk_size: int | None=None, output_buffer: torch.Tensor | None=None, retain_base_soma_prediction: bool=True) -> torch.Tensor", "parameters": [ { "name": "event_indices", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "event_values", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "state", "type": "AxoTemporalStreamingState", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "chunk_size", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "output_buffer", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "retain_base_soma_prediction", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "line": 771, "end_line": 921, "return_type": "torch.Tensor", "decorators": [] }, { "name": "recurrent_state_bytes", "kind": "method", "signature": "recurrent_state_bytes(self, *, dtype: torch.dtype) -> int", "parameters": [ { "name": "dtype", "type": "torch.dtype", "default": null, "kind": "keyword_only" } ], "line": 1127, "end_line": 1135, "return_type": "int", "decorators": [] }, { "name": "component_manifest", "kind": "method", "signature": "component_manifest(self) -> dict[str, object]", "parameters": [], "line": 1137, "end_line": 1163, "return_type": "dict[str, object]", "decorators": [] } ], "bases": [ "nn.Module" ] }, { "module": "temporal_core", "name": "create_temporal_model", "kind": "function", "page": "api/temporal-core", "line": 1166, "end_line": 1174, "signature": "create_temporal_model(source: BranchOfficialMamba, temporal_config: AxoTemporalCoreConfig) -> BranchOfficialMamba | AxoTemporalModel", "parameters": [ { "name": "source", "type": "BranchOfficialMamba", "default": null, "kind": "positional_or_keyword" }, { "name": "temporal_config", "type": "AxoTemporalCoreConfig", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "temporal_core", "name": "migrate_legacy_gru_state_dict", "kind": "function", "page": "api/temporal-core", "line": 1177, "end_line": 1195, "signature": "migrate_legacy_gru_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]", "parameters": [ { "name": "state_dict", "type": "dict[str, torch.Tensor]", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "block_forecast", "name": "BlockForecastConfig", "kind": "class", "page": "api/block-forecast", "line": 22, "end_line": 63, "signature": "BlockForecastConfig(core_kind: BlockForecastCore, patch_size: int, max_patch_rank: int = 8, max_trajectory_rank: int = 8, keep_local_tcn: bool = False, gru_hidden_units: int | None = None, residual_scale_init: float = 0.1, voltage_anchor_count: int | None = None, event_template: tuple[float, ...] | None = None, event_template_center_init: float = 0.9208316802978516, event_template_temperature: float = 0.25, event_template_hard: bool = False, behavior_adaptation: bool = False) -> None", "parameters": [ { "name": "core_kind", "type": "BlockForecastCore", "default": null, "kind": "positional_or_keyword" }, { "name": "patch_size", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "max_patch_rank", "type": "int", "default": "8", "kind": "positional_or_keyword" }, { "name": "max_trajectory_rank", "type": "int", "default": "8", "kind": "positional_or_keyword" }, { "name": "keep_local_tcn", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "gru_hidden_units", "type": "int | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "residual_scale_init", "type": "float", "default": "0.1", "kind": "positional_or_keyword" }, { "name": "voltage_anchor_count", "type": "int | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "event_template", "type": "tuple[float, ...] | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "event_template_center_init", "type": "float", "default": "0.9208316802978516", "kind": "positional_or_keyword" }, { "name": "event_template_temperature", "type": "float", "default": "0.25", "kind": "positional_or_keyword" }, { "name": "event_template_hard", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "behavior_adaptation", "type": "bool", "default": "False", "kind": "positional_or_keyword" } ], "fields": [ { "name": "core_kind", "type": "BlockForecastCore", "default": null }, { "name": "patch_size", "type": "int", "default": null }, { "name": "max_patch_rank", "type": "int", "default": "8" }, { "name": "max_trajectory_rank", "type": "int", "default": "8" }, { "name": "keep_local_tcn", "type": "bool", "default": "False" }, { "name": "gru_hidden_units", "type": "int | None", "default": "None" }, { "name": "residual_scale_init", "type": "float", "default": "0.1" }, { "name": "voltage_anchor_count", "type": "int | None", "default": "None" }, { "name": "event_template", "type": "tuple[float, ...] | None", "default": "None" }, { "name": "event_template_center_init", "type": "float", "default": "0.9208316802978516" }, { "name": "event_template_temperature", "type": "float", "default": "0.25" }, { "name": "event_template_hard", "type": "bool", "default": "False" }, { "name": "behavior_adaptation", "type": "bool", "default": "False" } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "block_forecast", "name": "CausalBlockForecastModel", "kind": "class", "page": "api/block-forecast", "line": 66, "end_line": 811, "signature": "CausalBlockForecastModel(self, source: nn.Module, forecast_config: BlockForecastConfig) -> None", "parameters": [ { "name": "source", "type": "nn.Module", "default": null, "kind": "positional_or_keyword" }, { "name": "forecast_config", "type": "BlockForecastConfig", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, source: nn.Module, forecast_config: BlockForecastConfig) -> None", "parameters": [ { "name": "source", "type": "nn.Module", "default": null, "kind": "positional_or_keyword" }, { "name": "forecast_config", "type": "BlockForecastConfig", "default": null, "kind": "positional_or_keyword" } ], "line": 69, "end_line": 206, "return_type": "None", "decorators": [] }, { "name": "config", "kind": "property", "signature": "CausalBlockForecastModel.config", "parameters": [], "line": 209, "end_line": 210, "return_type": null, "decorators": [ "property" ] }, { "name": "num_input", "kind": "property", "signature": "CausalBlockForecastModel.num_input: int", "parameters": [], "line": 213, "end_line": 214, "return_type": "int", "decorators": [ "property" ] }, { "name": "num_output", "kind": "property", "signature": "CausalBlockForecastModel.num_output: int", "parameters": [], "line": 217, "end_line": 218, "return_type": "int", "decorators": [ "property" ] }, { "name": "num_branch", "kind": "property", "signature": "CausalBlockForecastModel.num_branch: int", "parameters": [], "line": 221, "end_line": 222, "return_type": "int", "decorators": [ "property" ] }, { "name": "patch_size", "kind": "property", "signature": "CausalBlockForecastModel.patch_size: int", "parameters": [], "line": 225, "end_line": 226, "return_type": "int", "decorators": [ "property" ] }, { "name": "sparse_voltage_output", "kind": "property", "signature": "CausalBlockForecastModel.sparse_voltage_output: bool", "parameters": [], "line": 229, "end_line": 230, "return_type": "bool", "decorators": [ "property" ] }, { "name": "behavior_parameter_count", "kind": "property", "signature": "CausalBlockForecastModel.behavior_parameter_count: int", "parameters": [], "line": 233, "end_line": 238, "return_type": "int", "decorators": [ "property" ] }, { "name": "block_state_sequence", "kind": "method", "signature": "block_state_sequence(self, x: torch.Tensor, *, morphology_indices: torch.Tensor | None=None, behavior_parameters: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "x", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "behavior_parameters", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 447, "end_line": 485, "return_type": "torch.Tensor", "decorators": [] }, { "name": "decode_block_state_sequence", "kind": "method", "signature": "decode_block_state_sequence(self, block_state: torch.Tensor, *, behavior_parameters: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "block_state", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "behavior_parameters", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 487, "end_line": 509, "return_type": "torch.Tensor", "decorators": [] }, { "name": "forward", "kind": "method", "signature": "forward(self, x: torch.Tensor, *, morphology_indices: torch.Tensor | None=None, behavior_parameters: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "x", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "behavior_parameters", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 611, "end_line": 716, "return_type": "torch.Tensor", "decorators": [] }, { "name": "component_manifest", "kind": "method", "signature": "component_manifest(self) -> dict[str, object]", "parameters": [], "line": 782, "end_line": 811, "return_type": "dict[str, object]", "decorators": [] } ], "bases": [ "nn.Module" ] }, { "module": "support_surrogate", "name": "SupportP4Config", "kind": "class", "page": "api/lite", "line": 13, "end_line": 38, "signature": "SupportP4Config(input_dim: int = 1278, route_feature_dim: int = 118, token_dim: int = 58, state_dim: int = 16, patch_size: int = 4, output_dim: int = 2, morphology_ids: tuple[str, ...] = (), behavior_adaptation: bool = True) -> None", "parameters": [ { "name": "input_dim", "type": "int", "default": "1278", "kind": "positional_or_keyword" }, { "name": "route_feature_dim", "type": "int", "default": "118", "kind": "positional_or_keyword" }, { "name": "token_dim", "type": "int", "default": "58", "kind": "positional_or_keyword" }, { "name": "state_dim", "type": "int", "default": "16", "kind": "positional_or_keyword" }, { "name": "patch_size", "type": "int", "default": "4", "kind": "positional_or_keyword" }, { "name": "output_dim", "type": "int", "default": "2", "kind": "positional_or_keyword" }, { "name": "morphology_ids", "type": "tuple[str, ...]", "default": "()", "kind": "positional_or_keyword" }, { "name": "behavior_adaptation", "type": "bool", "default": "True", "kind": "positional_or_keyword" } ], "fields": [ { "name": "input_dim", "type": "int", "default": "1278" }, { "name": "route_feature_dim", "type": "int", "default": "118" }, { "name": "token_dim", "type": "int", "default": "58" }, { "name": "state_dim", "type": "int", "default": "16" }, { "name": "patch_size", "type": "int", "default": "4" }, { "name": "output_dim", "type": "int", "default": "2" }, { "name": "morphology_ids", "type": "tuple[str, ...]", "default": "()" }, { "name": "behavior_adaptation", "type": "bool", "default": "True" } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "support_surrogate", "name": "AdaptiveSupportP4Surrogate", "kind": "class", "page": "api/lite", "line": 41, "end_line": 535, "signature": "AdaptiveSupportP4Surrogate(self, config: SupportP4Config, route_features: torch.Tensor) -> None", "parameters": [ { "name": "config", "type": "SupportP4Config", "default": null, "kind": "positional_or_keyword" }, { "name": "route_features", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, config: SupportP4Config, route_features: torch.Tensor) -> None", "parameters": [ { "name": "config", "type": "SupportP4Config", "default": null, "kind": "positional_or_keyword" }, { "name": "route_features", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 44, "end_line": 115, "return_type": "None", "decorators": [] }, { "name": "patch_size", "kind": "property", "signature": "AdaptiveSupportP4Surrogate.patch_size: int", "parameters": [], "line": 118, "end_line": 119, "return_type": "int", "decorators": [ "property" ] }, { "name": "gate_feature_dim", "kind": "property", "signature": "AdaptiveSupportP4Surrogate.gate_feature_dim: int", "parameters": [], "line": 122, "end_line": 123, "return_type": "int", "decorators": [ "property" ] }, { "name": "reset_parameters", "kind": "method", "signature": "reset_parameters(self) -> None", "parameters": [], "line": 135, "end_line": 149, "return_type": "None", "decorators": [] }, { "name": "aggregate_inputs", "kind": "method", "signature": "aggregate_inputs(self, inputs: torch.Tensor, *, morphology_indices: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "inputs", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor", "default": null, "kind": "keyword_only" } ], "line": 151, "end_line": 172, "return_type": "torch.Tensor", "decorators": [] }, { "name": "compile_adaptation", "kind": "method", "signature": "compile_adaptation(self, behavior_parameters: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "behavior_parameters", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 174, "end_line": 186, "return_type": "torch.Tensor", "decorators": [] }, { "name": "initial_state", "kind": "method", "signature": "initial_state(self, batch_size: int, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor", "parameters": [ { "name": "batch_size", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "device", "type": "torch.device", "default": null, "kind": "keyword_only" }, { "name": "dtype", "type": "torch.dtype", "default": null, "kind": "keyword_only" } ], "line": 188, "end_line": 200, "return_type": "torch.Tensor", "decorators": [] }, { "name": "runtime_coefficient_slices", "kind": "property", "signature": "AdaptiveSupportP4Surrogate.runtime_coefficient_slices: dict[str, slice]", "parameters": [], "line": 212, "end_line": 248, "return_type": "dict[str, slice]", "decorators": [ "property" ] }, { "name": "step_token", "kind": "method", "signature": "step_token(self, token: torch.Tensor, state: torch.Tensor, *, adaptation_cache: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]", "parameters": [ { "name": "token", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "state", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "adaptation_cache", "type": "torch.Tensor | None", "default": null, "kind": "keyword_only" } ], "line": 250, "end_line": 312, "return_type": "tuple[torch.Tensor, torch.Tensor]", "decorators": [] }, { "name": "step_p4", "kind": "method", "signature": "step_p4(self, feature_patch: torch.Tensor, state: torch.Tensor, *, behavior_parameters: torch.Tensor | None=None, adaptation_cache: torch.Tensor | None=None) -> tuple[torch.Tensor, torch.Tensor]", "parameters": [ { "name": "feature_patch", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "state", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "behavior_parameters", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "adaptation_cache", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 314, "end_line": 342, "return_type": "tuple[torch.Tensor, torch.Tensor]", "decorators": [] }, { "name": "forward_feature_summaries", "kind": "method", "signature": "forward_feature_summaries(self, summaries: torch.Tensor, *, morphology_indices: torch.Tensor | None=None, behavior_parameters: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "summaries", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "behavior_parameters", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 453, "end_line": 469, "return_type": "torch.Tensor", "decorators": [] }, { "name": "forward_runtime_coefficients", "kind": "method", "signature": "forward_runtime_coefficients(self, summaries: torch.Tensor, *, runtime_coefficients: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "summaries", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "runtime_coefficients", "type": "torch.Tensor", "default": null, "kind": "keyword_only" } ], "line": 471, "end_line": 489, "return_type": "torch.Tensor", "decorators": [] }, { "name": "forward_with_gate_features", "kind": "method", "signature": "forward_with_gate_features(self, inputs: torch.Tensor, *, morphology_indices: torch.Tensor | None=None, behavior_parameters: torch.Tensor | None=None) -> tuple[torch.Tensor, torch.Tensor]", "parameters": [ { "name": "inputs", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "behavior_parameters", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 491, "end_line": 516, "return_type": "tuple[torch.Tensor, torch.Tensor]", "decorators": [] }, { "name": "forward", "kind": "method", "signature": "forward(self, inputs: torch.Tensor, *, morphology_indices: torch.Tensor | None=None, behavior_parameters: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "inputs", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "behavior_parameters", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 518, "end_line": 535, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [ "nn.Module" ] }, { "module": "model", "name": "BranchELMConfig", "kind": "class", "page": "api/model", "line": 10, "end_line": 20, "signature": "BranchELMConfig(input_dim: int = 1278, memory_units: int = 30, num_branches: int = 32, hidden_units: int = 64, synapse_decay: float = 0.85, memory_decay: float = 0.9) -> None", "parameters": [ { "name": "input_dim", "type": "int", "default": "1278", "kind": "positional_or_keyword" }, { "name": "memory_units", "type": "int", "default": "30", "kind": "positional_or_keyword" }, { "name": "num_branches", "type": "int", "default": "32", "kind": "positional_or_keyword" }, { "name": "hidden_units", "type": "int", "default": "64", "kind": "positional_or_keyword" }, { "name": "synapse_decay", "type": "float", "default": "0.85", "kind": "positional_or_keyword" }, { "name": "memory_decay", "type": "float", "default": "0.9", "kind": "positional_or_keyword" } ], "fields": [ { "name": "input_dim", "type": "int", "default": "1278" }, { "name": "memory_units", "type": "int", "default": "30" }, { "name": "num_branches", "type": "int", "default": "32" }, { "name": "hidden_units", "type": "int", "default": "64" }, { "name": "synapse_decay", "type": "float", "default": "0.85" }, { "name": "memory_decay", "type": "float", "default": "0.9" } ], "methods": [ { "name": "branch_elm_30", "kind": "classmethod", "signature": "branch_elm_30(cls, *, input_dim: int=1278, num_branches: int=32) -> 'BranchELMConfig'", "parameters": [ { "name": "input_dim", "type": "int", "default": "1278", "kind": "keyword_only", "resolved_default": 1278 }, { "name": "num_branches", "type": "int", "default": "32", "kind": "keyword_only", "resolved_default": 32 } ], "line": 19, "end_line": 20, "return_type": "'BranchELMConfig'", "decorators": [ "classmethod" ] } ], "bases": [], "inherited_fields": [] }, { "module": "model", "name": "BranchELM", "kind": "class", "page": "api/model", "line": 23, "end_line": 74, "signature": "BranchELM(self, config: BranchELMConfig) -> None", "parameters": [ { "name": "config", "type": "BranchELMConfig", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, config: BranchELMConfig) -> None", "parameters": [ { "name": "config", "type": "BranchELMConfig", "default": null, "kind": "positional_or_keyword" } ], "line": 26, "end_line": 46, "return_type": "None", "decorators": [] }, { "name": "forward", "kind": "method", "signature": "forward(self, x: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "x", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 48, "end_line": 71, "return_type": "torch.Tensor", "decorators": [] }, { "name": "parameter_count", "kind": "method", "signature": "parameter_count(self) -> int", "parameters": [], "line": 73, "end_line": 74, "return_type": "int", "decorators": [] } ], "bases": [ "nn.Module" ] }, { "module": "checkpoint", "name": "save_checkpoint", "kind": "function", "page": "api/checkpoint", "line": 188, "end_line": 206, "signature": "save_checkpoint(model: torch.nn.Module, path: str | Path, *, metadata: dict[str, Any] | None=None) -> None", "parameters": [ { "name": "model", "type": "torch.nn.Module", "default": null, "kind": "positional_or_keyword" }, { "name": "path", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "metadata", "type": "dict[str, Any] | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "fields": [], "methods": [] }, { "module": "checkpoint", "name": "load_checkpoint", "kind": "function", "page": "api/checkpoint", "line": 209, "end_line": 329, "signature": "load_checkpoint(path: str | Path, *, map_location: str='cpu') -> tuple[torch.nn.Module, dict[str, Any]]", "parameters": [ { "name": "path", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "map_location", "type": "str", "default": "'cpu'", "kind": "keyword_only", "resolved_default": "cpu" } ], "fields": [], "methods": [] }, { "module": "adaptation", "name": "CudaGraphAdaptationStep", "kind": "class", "page": "api/adaptation", "line": 55, "end_line": 169, "signature": "CudaGraphAdaptationStep(module: nn.Module, optimizer: torch.optim.Optimizer, static_inputs: torch.Tensor, static_targets: torch.Tensor, graph: torch.cuda.CUDAGraph, static_loss: torch.Tensor) -> None", "parameters": [ { "name": "module", "type": "nn.Module", "default": null, "kind": "positional_or_keyword" }, { "name": "optimizer", "type": "torch.optim.Optimizer", "default": null, "kind": "positional_or_keyword" }, { "name": "static_inputs", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "static_targets", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "graph", "type": "torch.cuda.CUDAGraph", "default": null, "kind": "positional_or_keyword" }, { "name": "static_loss", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "module", "type": "nn.Module", "default": null }, { "name": "optimizer", "type": "torch.optim.Optimizer", "default": null }, { "name": "static_inputs", "type": "torch.Tensor", "default": null }, { "name": "static_targets", "type": "torch.Tensor", "default": null }, { "name": "graph", "type": "torch.cuda.CUDAGraph", "default": null }, { "name": "static_loss", "type": "torch.Tensor", "default": null } ], "methods": [ { "name": "capture", "kind": "classmethod", "signature": "capture(cls, module: nn.Module, optimizer: torch.optim.Optimizer, example_inputs: torch.Tensor, example_targets: torch.Tensor, loss_function: TensorLoss, *, warmup_steps: int=3, gradient_transform: GradientTransform | None=None) -> 'CudaGraphAdaptationStep'", "parameters": [ { "name": "module", "type": "nn.Module", "default": null, "kind": "positional_or_keyword" }, { "name": "optimizer", "type": "torch.optim.Optimizer", "default": null, "kind": "positional_or_keyword" }, { "name": "example_inputs", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "example_targets", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "loss_function", "type": "TensorLoss", "default": null, "kind": "positional_or_keyword" }, { "name": "warmup_steps", "type": "int", "default": "3", "kind": "keyword_only", "resolved_default": 3 }, { "name": "gradient_transform", "type": "GradientTransform | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 66, "end_line": 141, "return_type": "'CudaGraphAdaptationStep'", "decorators": [ "classmethod" ] }, { "name": "__call__", "kind": "method", "signature": "__call__(self, inputs: torch.Tensor, targets: torch.Tensor, *, synchronize: bool=False) -> torch.Tensor", "parameters": [ { "name": "inputs", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "targets", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "synchronize", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false } ], "line": 143, "end_line": 169, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [], "inherited_fields": [] }, { "module": "simulation_contract", "name": "LargePopulationSimulationContract", "kind": "class", "page": "api/simulation-contract", "line": 9, "end_line": 272, "signature": "LargePopulationSimulationContract(name: str, population_size: int, step_ms: float, input_step_ms: float, output_step_ms: float, target_latency_ms: float, model_family: str, patch_size: int, temporal_execution_mode: str, patch_phase_mode: str, patch_phase_assignment: str, patch_phase_storage_layout: str, lossless_ordered_temporal_io: bool, causal_block_feedback_required: bool, minimum_routing_delay_steps: int, synaptic_efficacy_values_per_neuron: int, adaptation_storage_bits: int, activation_storage_bits: int, recurrent_state_storage_bits: int, external_drive_storage_bits: int, morphology_classes: int, firing_rate: float, fanout: int, recurrent_contact_fraction: float, excitatory_fraction: float, local_connection_fraction: float, spatial_tile_neurons: int, delay_slots: int, branch_count: int, local_source_identity_required: bool, long_range_source_identity_required: bool, local_branch_input_mode: str, input_channel_collision_mode: str, local_source_summary_is_lossless_for_declared_router: bool, posthoc_source_projection_allowed: bool, morphology_assignment_independent: bool, activity_must_be_model_generated: bool, quota_activity_forcing_allowed: bool, activity_selection_mode: str, adaptation_reference_envelope_validation_required: bool, morphology_aware_branch_targeting: bool = True, source_type_aware_branch_targeting: bool = True, source_discovery_included: bool = True, event_delivery_included: bool = True, delay_queue_included: bool = True, every_neuron_has_custom_adaptation: bool = True, empirical_connectome: bool = False) -> None", "parameters": [ { "name": "name", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "population_size", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "step_ms", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "input_step_ms", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "output_step_ms", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "target_latency_ms", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "model_family", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "patch_size", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "temporal_execution_mode", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "patch_phase_mode", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "patch_phase_assignment", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "patch_phase_storage_layout", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "lossless_ordered_temporal_io", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "causal_block_feedback_required", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "minimum_routing_delay_steps", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "synaptic_efficacy_values_per_neuron", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "adaptation_storage_bits", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "activation_storage_bits", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "recurrent_state_storage_bits", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "external_drive_storage_bits", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_classes", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "firing_rate", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "fanout", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "recurrent_contact_fraction", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "excitatory_fraction", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "local_connection_fraction", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "spatial_tile_neurons", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "delay_slots", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "branch_count", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "local_source_identity_required", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "long_range_source_identity_required", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "local_branch_input_mode", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "input_channel_collision_mode", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "local_source_summary_is_lossless_for_declared_router", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "posthoc_source_projection_allowed", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_assignment_independent", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "activity_must_be_model_generated", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "quota_activity_forcing_allowed", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "activity_selection_mode", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "adaptation_reference_envelope_validation_required", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_aware_branch_targeting", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "source_type_aware_branch_targeting", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "source_discovery_included", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "event_delivery_included", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "delay_queue_included", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "every_neuron_has_custom_adaptation", "type": "bool", "default": "True", "kind": "positional_or_keyword" }, { "name": "empirical_connectome", "type": "bool", "default": "False", "kind": "positional_or_keyword" } ], "fields": [ { "name": "name", "type": "str", "default": null }, { "name": "population_size", "type": "int", "default": null }, { "name": "step_ms", "type": "float", "default": null }, { "name": "input_step_ms", "type": "float", "default": null }, { "name": "output_step_ms", "type": "float", "default": null }, { "name": "target_latency_ms", "type": "float", "default": null }, { "name": "model_family", "type": "str", "default": null }, { "name": "patch_size", "type": "int", "default": null }, { "name": "temporal_execution_mode", "type": "str", "default": null }, { "name": "patch_phase_mode", "type": "str", "default": null }, { "name": "patch_phase_assignment", "type": "str", "default": null }, { "name": "patch_phase_storage_layout", "type": "str", "default": null }, { "name": "lossless_ordered_temporal_io", "type": "bool", "default": null }, { "name": "causal_block_feedback_required", "type": "bool", "default": null }, { "name": "minimum_routing_delay_steps", "type": "int", "default": null }, { "name": "synaptic_efficacy_values_per_neuron", "type": "int", "default": null }, { "name": "adaptation_storage_bits", "type": "int", "default": null }, { "name": "activation_storage_bits", "type": "int", "default": null }, { "name": "recurrent_state_storage_bits", "type": "int", "default": null }, { "name": "external_drive_storage_bits", "type": "int", "default": null }, { "name": "morphology_classes", "type": "int", "default": null }, { "name": "firing_rate", "type": "float", "default": null }, { "name": "fanout", "type": "int", "default": null }, { "name": "recurrent_contact_fraction", "type": "float", "default": null }, { "name": "excitatory_fraction", "type": "float", "default": null }, { "name": "local_connection_fraction", "type": "float", "default": null }, { "name": "spatial_tile_neurons", "type": "int", "default": null }, { "name": "delay_slots", "type": "int", "default": null }, { "name": "branch_count", "type": "int", "default": null }, { "name": "local_source_identity_required", "type": "bool", "default": null }, { "name": "long_range_source_identity_required", "type": "bool", "default": null }, { "name": "local_branch_input_mode", "type": "str", "default": null }, { "name": "input_channel_collision_mode", "type": "str", "default": null }, { "name": "local_source_summary_is_lossless_for_declared_router", "type": "bool", "default": null }, { "name": "posthoc_source_projection_allowed", "type": "bool", "default": null }, { "name": "morphology_assignment_independent", "type": "bool", "default": null }, { "name": "activity_must_be_model_generated", "type": "bool", "default": null }, { "name": "quota_activity_forcing_allowed", "type": "bool", "default": null }, { "name": "activity_selection_mode", "type": "str", "default": null }, { "name": "adaptation_reference_envelope_validation_required", "type": "bool", "default": null }, { "name": "morphology_aware_branch_targeting", "type": "bool", "default": "True" }, { "name": "source_type_aware_branch_targeting", "type": "bool", "default": "True" }, { "name": "source_discovery_included", "type": "bool", "default": "True" }, { "name": "event_delivery_included", "type": "bool", "default": "True" }, { "name": "delay_queue_included", "type": "bool", "default": "True" }, { "name": "every_neuron_has_custom_adaptation", "type": "bool", "default": "True" }, { "name": "empirical_connectome", "type": "bool", "default": "False" } ], "methods": [ { "name": "steps_per_second", "kind": "property", "signature": "LargePopulationSimulationContract.steps_per_second: float", "parameters": [], "line": 216, "end_line": 217, "return_type": "float", "decorators": [ "property" ] }, { "name": "target_neuron_steps_per_second", "kind": "property", "signature": "LargePopulationSimulationContract.target_neuron_steps_per_second: float", "parameters": [], "line": 220, "end_line": 221, "return_type": "float", "decorators": [ "property" ] }, { "name": "largest_patch_phase_population", "kind": "property", "signature": "LargePopulationSimulationContract.largest_patch_phase_population: int", "parameters": [], "line": 224, "end_line": 227, "return_type": "int", "decorators": [ "property" ] }, { "name": "active_sources_per_step", "kind": "property", "signature": "LargePopulationSimulationContract.active_sources_per_step: int", "parameters": [], "line": 230, "end_line": 231, "return_type": "int", "decorators": [ "property" ] }, { "name": "events_per_step", "kind": "property", "signature": "LargePopulationSimulationContract.events_per_step: int", "parameters": [], "line": 234, "end_line": 239, "return_type": "int", "decorators": [ "property" ] }, { "name": "synaptic_efficacy_bank_bytes", "kind": "property", "signature": "LargePopulationSimulationContract.synaptic_efficacy_bank_bytes: int", "parameters": [], "line": 242, "end_line": 248, "return_type": "int", "decorators": [ "property" ] }, { "name": "delay_queue_bytes", "kind": "property", "signature": "LargePopulationSimulationContract.delay_queue_bytes: int", "parameters": [], "line": 251, "end_line": 258, "return_type": "int", "decorators": [ "property" ] }, { "name": "external_drive_block_bytes", "kind": "property", "signature": "LargePopulationSimulationContract.external_drive_block_bytes: int", "parameters": [], "line": 261, "end_line": 268, "return_type": "int", "decorators": [ "property" ] }, { "name": "bytes_to_gib", "kind": "staticmethod", "signature": "bytes_to_gib(byte_count: int) -> float", "parameters": [ { "name": "byte_count", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "line": 271, "end_line": 272, "return_type": "float", "decorators": [ "staticmethod" ] } ], "bases": [], "inherited_fields": [] }, { "module": "simulation_contract", "name": "PopulationInferenceProfile", "kind": "class", "page": "api/simulation-contract", "line": 276, "end_line": 361, "signature": "PopulationInferenceProfile(name: str, support_model_id: str, support_token_width: int, support_state_width: int, support_block_rows: int, support_block_token: int, support_block_state: int, external_block_width: int, project_block_width: int, compiled_adaptation_values_per_neuron: int, morphology_adapter_values_per_class: int, synaptic_efficacy_values_per_neuron: int, output_export_fraction: float, primary_metric: str, minimum_neuron_steps_per_second: float, measured_median_neuron_steps_per_second: float, measured_p90_neuron_steps_per_second: float, measured_voltage_sera_mv2: float, measured_voltage_root_sera_mv: float, measured_spike_mean_f1: float, measured_population_firing_rate_hz: float, population_adaptation_initialization: str, population_adaptation_individually_fitted: bool, individual_behavior_adaptation_supported: bool, individual_synaptic_adaptation_supported: bool) -> None", "parameters": [ { "name": "name", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "support_model_id", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "support_token_width", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "support_state_width", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "support_block_rows", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "support_block_token", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "support_block_state", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "external_block_width", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "project_block_width", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "compiled_adaptation_values_per_neuron", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_adapter_values_per_class", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "synaptic_efficacy_values_per_neuron", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "output_export_fraction", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "primary_metric", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "minimum_neuron_steps_per_second", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "measured_median_neuron_steps_per_second", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "measured_p90_neuron_steps_per_second", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "measured_voltage_sera_mv2", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "measured_voltage_root_sera_mv", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "measured_spike_mean_f1", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "measured_population_firing_rate_hz", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "population_adaptation_initialization", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "population_adaptation_individually_fitted", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "individual_behavior_adaptation_supported", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "individual_synaptic_adaptation_supported", "type": "bool", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "name", "type": "str", "default": null }, { "name": "support_model_id", "type": "str", "default": null }, { "name": "support_token_width", "type": "int", "default": null }, { "name": "support_state_width", "type": "int", "default": null }, { "name": "support_block_rows", "type": "int", "default": null }, { "name": "support_block_token", "type": "int", "default": null }, { "name": "support_block_state", "type": "int", "default": null }, { "name": "external_block_width", "type": "int", "default": null }, { "name": "project_block_width", "type": "int", "default": null }, { "name": "compiled_adaptation_values_per_neuron", "type": "int", "default": null }, { "name": "morphology_adapter_values_per_class", "type": "int", "default": null }, { "name": "synaptic_efficacy_values_per_neuron", "type": "int", "default": null }, { "name": "output_export_fraction", "type": "float", "default": null }, { "name": "primary_metric", "type": "str", "default": null }, { "name": "minimum_neuron_steps_per_second", "type": "float", "default": null }, { "name": "measured_median_neuron_steps_per_second", "type": "float", "default": null }, { "name": "measured_p90_neuron_steps_per_second", "type": "float", "default": null }, { "name": "measured_voltage_sera_mv2", "type": "float", "default": null }, { "name": "measured_voltage_root_sera_mv", "type": "float", "default": null }, { "name": "measured_spike_mean_f1", "type": "float", "default": null }, { "name": "measured_population_firing_rate_hz", "type": "float", "default": null }, { "name": "population_adaptation_initialization", "type": "str", "default": null }, { "name": "population_adaptation_individually_fitted", "type": "bool", "default": null }, { "name": "individual_behavior_adaptation_supported", "type": "bool", "default": null }, { "name": "individual_synaptic_adaptation_supported", "type": "bool", "default": null } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "simulation_contract", "name": "MILLION_NEURON_REALTIME_CONTRACT", "kind": "constant", "page": "api/simulation-contract", "line": 362, "end_line": 405, "signature": "MILLION_NEURON_REALTIME_CONTRACT = LargePopulationSimulationContract(name='million-neuron-realtime-rtx5090', population_size=1048576, step_ms=1.0, input_step_ms=1.0, output_step_ms=1.0, target_latency_ms=1.0, model_family='axosim-lite-causal-block-gru-regression', patch_size=4, temporal_execution_mode='causal_next_block_forecast', patch_phase_mode='staggered', patch_phase_assignment='morphology_role_stratified_random', patch_phase_storage_layout='phase_tile_major', lossless_ordered_temporal_io=True, causal_block_feedback_required=True, minimum_routing_delay_steps=4, synaptic_efficacy_values_per_neuron=362, adaptation_storage_bits=8, activation_storage_bits=16, recurrent_state_storage_bits=16, external_drive_storage_bits=8, morphology_classes=5, firing_rate=0.003869090909090909, fanout=1448, recurrent_contact_fraction=1.0 / 4.0, excitatory_fraction=0.8, local_connection_fraction=0.9, spatial_tile_neurons=4096, delay_slots=8, branch_count=118, local_source_identity_required=True, long_range_source_identity_required=True, local_branch_input_mode='binary_channels_then_exact_weighted_branch_currents', input_channel_collision_mode='binary_or_before_branch_weighting', local_source_summary_is_lossless_for_declared_router=True, posthoc_source_projection_allowed=False, morphology_assignment_independent=True, activity_must_be_model_generated=True, quota_activity_forcing_allowed=False, activity_selection_mode='fixed_model_score_threshold_crossing', adaptation_reference_envelope_validation_required=True)", "parameters": [], "fields": [], "methods": [] }, { "module": "simulation_contract", "name": "AXOSIM_POPULATION_PROFILE", "kind": "constant", "page": "api/simulation-contract", "line": 407, "end_line": 437, "signature": "AXOSIM_POPULATION_PROFILE = PopulationInferenceProfile(name='axosim-population', support_model_id='support_t24_s8_distilled_i600', support_token_width=24, support_state_width=8, support_block_rows=32, support_block_token=32, support_block_state=16, external_block_width=32, project_block_width=64, compiled_adaptation_values_per_neuron=88, morphology_adapter_values_per_class=88, synaptic_efficacy_values_per_neuron=362, output_export_fraction=0.01, primary_metric='neuron_steps_per_second', minimum_neuron_steps_per_second=MILLION_NEURON_REALTIME_CONTRACT.target_neuron_steps_per_second, measured_median_neuron_steps_per_second=1201327144.0976055, measured_p90_neuron_steps_per_second=1177454556.7997348, measured_voltage_sera_mv2=225.82762062058336, measured_voltage_root_sera_mv=15.02756203183282, measured_spike_mean_f1=0.35383167978104685, measured_population_firing_rate_hz=3.956115245819092, population_adaptation_initialization='five_fitted_morphology_rows_plus_direct_w8_behavior_residuals', population_adaptation_individually_fitted=False, individual_behavior_adaptation_supported=True, individual_synaptic_adaptation_supported=True)", "parameters": [], "fields": [], "methods": [] }, { "module": "population", "name": "QuantizedNeuronBehaviorBank", "kind": "class", "page": "api/population", "line": 20, "end_line": 33, "signature": "QuantizedNeuronBehaviorBank(values: torch.Tensor, scales: torch.Tensor, bits: int, logical_parameters_per_neuron: int) -> None", "parameters": [ { "name": "values", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "scales", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "bits", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "logical_parameters_per_neuron", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "values", "type": "torch.Tensor", "default": null }, { "name": "scales", "type": "torch.Tensor", "default": null }, { "name": "bits", "type": "int", "default": null }, { "name": "logical_parameters_per_neuron", "type": "int", "default": null } ], "methods": [ { "name": "storage_bytes", "kind": "property", "signature": "QuantizedNeuronBehaviorBank.storage_bytes: int", "parameters": [], "line": 29, "end_line": 33, "return_type": "int", "decorators": [ "property" ] } ], "bases": [], "inherited_fields": [] }, { "module": "population", "name": "MixedQuantizedNeuronBehaviorBank", "kind": "class", "page": "api/population", "line": 37, "end_line": 51, "signature": "MixedQuantizedNeuronBehaviorBank(values: torch.Tensor, scales: torch.Tensor, byte_slices: dict[str, slice], component_bits: dict[str, int], logical_parameters_per_neuron: int) -> None", "parameters": [ { "name": "values", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "scales", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "byte_slices", "type": "dict[str, slice]", "default": null, "kind": "positional_or_keyword" }, { "name": "component_bits", "type": "dict[str, int]", "default": null, "kind": "positional_or_keyword" }, { "name": "logical_parameters_per_neuron", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "values", "type": "torch.Tensor", "default": null }, { "name": "scales", "type": "torch.Tensor", "default": null }, { "name": "byte_slices", "type": "dict[str, slice]", "default": null }, { "name": "component_bits", "type": "dict[str, int]", "default": null }, { "name": "logical_parameters_per_neuron", "type": "int", "default": null } ], "methods": [ { "name": "storage_bytes", "kind": "property", "signature": "MixedQuantizedNeuronBehaviorBank.storage_bytes: int", "parameters": [], "line": 47, "end_line": 51, "return_type": "int", "decorators": [ "property" ] } ], "bases": [], "inherited_fields": [] }, { "module": "population", "name": "quantize_neuron_behavior_parameters", "kind": "function", "page": "api/population", "line": 54, "end_line": 108, "signature": "quantize_neuron_behavior_parameters(parameters: torch.Tensor, *, adaptation: NeuronBehaviorAdaptation, bits: int) -> QuantizedNeuronBehaviorBank", "parameters": [ { "name": "parameters", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "adaptation", "type": "NeuronBehaviorAdaptation", "default": null, "kind": "keyword_only" }, { "name": "bits", "type": "int", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [] }, { "module": "population", "name": "quantize_mixed_neuron_behavior_parameters", "kind": "function", "page": "api/population", "line": 111, "end_line": 183, "signature": "quantize_mixed_neuron_behavior_parameters(parameters: torch.Tensor, *, adaptation: NeuronBehaviorAdaptation, component_bits: dict[str, int]) -> MixedQuantizedNeuronBehaviorBank", "parameters": [ { "name": "parameters", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "adaptation", "type": "NeuronBehaviorAdaptation", "default": null, "kind": "keyword_only" }, { "name": "component_bits", "type": "dict[str, int]", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [] }, { "module": "population", "name": "dequantize_neuron_behavior_parameters", "kind": "function", "page": "api/population", "line": 186, "end_line": 215, "signature": "dequantize_neuron_behavior_parameters(bank: QuantizedNeuronBehaviorBank, *, adaptation: NeuronBehaviorAdaptation) -> torch.Tensor", "parameters": [ { "name": "bank", "type": "QuantizedNeuronBehaviorBank", "default": null, "kind": "positional_or_keyword" }, { "name": "adaptation", "type": "NeuronBehaviorAdaptation", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [] }, { "module": "population", "name": "dequantize_mixed_neuron_behavior_parameters", "kind": "function", "page": "api/population", "line": 218, "end_line": 251, "signature": "dequantize_mixed_neuron_behavior_parameters(bank: MixedQuantizedNeuronBehaviorBank, *, adaptation: NeuronBehaviorAdaptation) -> torch.Tensor", "parameters": [ { "name": "bank", "type": "MixedQuantizedNeuronBehaviorBank", "default": null, "kind": "positional_or_keyword" }, { "name": "adaptation", "type": "NeuronBehaviorAdaptation", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [] }, { "module": "population", "name": "pack_neuron_behavior_adapter", "kind": "function", "page": "api/population", "line": 341, "end_line": 393, "signature": "pack_neuron_behavior_adapter(adapter: NeuronBehaviorAdapter, *, adaptation: NeuronBehaviorAdaptation) -> torch.Tensor", "parameters": [ { "name": "adapter", "type": "NeuronBehaviorAdapter", "default": null, "kind": "positional_or_keyword" }, { "name": "adaptation", "type": "NeuronBehaviorAdaptation", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [] }, { "module": "population", "name": "GRUPopulationRunner", "kind": "class", "page": "api/population", "line": 1344, "end_line": 1556, "signature": "GRUPopulationRunner(self, model: AxoTemporalModel, *, population: int, chunk_size: int, compile_step: bool=True, compile_mode: str='reduce-overhead', population_embedding_adapters: torch.Tensor | None=None) -> None", "parameters": [ { "name": "model", "type": "AxoTemporalModel", "default": null, "kind": "positional_or_keyword" }, { "name": "population", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "chunk_size", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "compile_step", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "compile_mode", "type": "str", "default": "'reduce-overhead'", "kind": "keyword_only", "resolved_default": "reduce-overhead" }, { "name": "population_embedding_adapters", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, model: AxoTemporalModel, *, population: int, chunk_size: int, compile_step: bool=True, compile_mode: str='reduce-overhead', population_embedding_adapters: torch.Tensor | None=None) -> None", "parameters": [ { "name": "model", "type": "AxoTemporalModel", "default": null, "kind": "positional_or_keyword" }, { "name": "population", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "chunk_size", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "compile_step", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "compile_mode", "type": "str", "default": "'reduce-overhead'", "kind": "keyword_only", "resolved_default": "reduce-overhead" }, { "name": "population_embedding_adapters", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 1347, "end_line": 1441, "return_type": "None", "decorators": [] }, { "name": "recurrent_state", "kind": "property", "signature": "GRUPopulationRunner.recurrent_state: torch.Tensor", "parameters": [], "line": 1444, "end_line": 1447, "return_type": "torch.Tensor", "decorators": [ "property" ] }, { "name": "patch_sum", "kind": "property", "signature": "GRUPopulationRunner.patch_sum: torch.Tensor | None", "parameters": [], "line": 1450, "end_line": 1451, "return_type": "torch.Tensor | None", "decorators": [ "property" ] }, { "name": "patch_correction", "kind": "property", "signature": "GRUPopulationRunner.patch_correction: torch.Tensor | None", "parameters": [], "line": 1454, "end_line": 1455, "return_type": "torch.Tensor | None", "decorators": [ "property" ] }, { "name": "patch_position", "kind": "property", "signature": "GRUPopulationRunner.patch_position: int", "parameters": [], "line": 1458, "end_line": 1459, "return_type": "int", "decorators": [ "property" ] }, { "name": "step", "kind": "method", "signature": "step(self, event_indices: torch.Tensor, event_values: torch.Tensor, morphology_indices: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "event_indices", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "event_values", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 1461, "end_line": 1526, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [] }, { "module": "population", "name": "GRUBranchPopulationRunner", "kind": "class", "page": "api/population", "line": 1559, "end_line": 1701, "signature": "GRUBranchPopulationRunner(self, model: AxoTemporalModel, *, population: int, chunk_size: int, compile_step: bool=True, compile_mode: str='reduce-overhead', population_embedding_adapters: torch.Tensor | None=None) -> None", "parameters": [ { "name": "model", "type": "AxoTemporalModel", "default": null, "kind": "positional_or_keyword" }, { "name": "population", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "chunk_size", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "compile_step", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "compile_mode", "type": "str", "default": "'reduce-overhead'", "kind": "keyword_only", "resolved_default": "reduce-overhead" }, { "name": "population_embedding_adapters", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, model: AxoTemporalModel, *, population: int, chunk_size: int, compile_step: bool=True, compile_mode: str='reduce-overhead', population_embedding_adapters: torch.Tensor | None=None) -> None", "parameters": [ { "name": "model", "type": "AxoTemporalModel", "default": null, "kind": "positional_or_keyword" }, { "name": "population", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "chunk_size", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "compile_step", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "compile_mode", "type": "str", "default": "'reduce-overhead'", "kind": "keyword_only", "resolved_default": "reduce-overhead" }, { "name": "population_embedding_adapters", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 1562, "end_line": 1617, "return_type": "None", "decorators": [] }, { "name": "step", "kind": "method", "signature": "step(self, branch_inputs: torch.Tensor, morphology_indices: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "branch_inputs", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 1619, "end_line": 1701, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [ "GRUPopulationRunner" ] }, { "module": "population", "name": "AdaptedGRUBranchPopulationRunner", "kind": "class", "page": "api/population", "line": 1704, "end_line": 1995, "signature": "AdaptedGRUBranchPopulationRunner(self, model: AxoTemporalModel, *, population: int, chunk_size: int, behavior_parameters: torch.Tensor | QuantizedNeuronBehaviorBank | MixedQuantizedNeuronBehaviorBank, adaptation: NeuronBehaviorAdaptation, compile_step: bool=True, compile_mode: str='max-autotune-no-cudagraphs') -> None", "parameters": [ { "name": "model", "type": "AxoTemporalModel", "default": null, "kind": "positional_or_keyword" }, { "name": "population", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "chunk_size", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "behavior_parameters", "type": "torch.Tensor | QuantizedNeuronBehaviorBank | MixedQuantizedNeuronBehaviorBank", "default": null, "kind": "keyword_only" }, { "name": "adaptation", "type": "NeuronBehaviorAdaptation", "default": null, "kind": "keyword_only" }, { "name": "compile_step", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "compile_mode", "type": "str", "default": "'max-autotune-no-cudagraphs'", "kind": "keyword_only", "resolved_default": "max-autotune-no-cudagraphs" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, model: AxoTemporalModel, *, population: int, chunk_size: int, behavior_parameters: torch.Tensor | QuantizedNeuronBehaviorBank | MixedQuantizedNeuronBehaviorBank, adaptation: NeuronBehaviorAdaptation, compile_step: bool=True, compile_mode: str='max-autotune-no-cudagraphs') -> None", "parameters": [ { "name": "model", "type": "AxoTemporalModel", "default": null, "kind": "positional_or_keyword" }, { "name": "population", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "chunk_size", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "behavior_parameters", "type": "torch.Tensor | QuantizedNeuronBehaviorBank | MixedQuantizedNeuronBehaviorBank", "default": null, "kind": "keyword_only" }, { "name": "adaptation", "type": "NeuronBehaviorAdaptation", "default": null, "kind": "keyword_only" }, { "name": "compile_step", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "compile_mode", "type": "str", "default": "'max-autotune-no-cudagraphs'", "kind": "keyword_only", "resolved_default": "max-autotune-no-cudagraphs" } ], "line": 1709, "end_line": 1918, "return_type": "None", "decorators": [] }, { "name": "step", "kind": "method", "signature": "step(self, branch_inputs: torch.Tensor, morphology_indices: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "branch_inputs", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 1920, "end_line": 1995, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [ "GRUBranchPopulationRunner" ] }, { "module": "population", "name": "GroupedGRUPopulationRunner", "kind": "class", "page": "api/population", "line": 1998, "end_line": 2270, "signature": "GroupedGRUPopulationRunner(self, model: AxoTemporalModel, *, population: int, groups: int, chunk_size: int, compile_step: bool=True, compile_mode: str='reduce-overhead') -> None", "parameters": [ { "name": "model", "type": "AxoTemporalModel", "default": null, "kind": "positional_or_keyword" }, { "name": "population", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "groups", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "chunk_size", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "compile_step", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "compile_mode", "type": "str", "default": "'reduce-overhead'", "kind": "keyword_only", "resolved_default": "reduce-overhead" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, model: AxoTemporalModel, *, population: int, groups: int, chunk_size: int, compile_step: bool=True, compile_mode: str='reduce-overhead') -> None", "parameters": [ { "name": "model", "type": "AxoTemporalModel", "default": null, "kind": "positional_or_keyword" }, { "name": "population", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "groups", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "chunk_size", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "compile_step", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "compile_mode", "type": "str", "default": "'reduce-overhead'", "kind": "keyword_only", "resolved_default": "reduce-overhead" } ], "line": 2001, "end_line": 2109, "return_type": "None", "decorators": [] }, { "name": "resident_parameter_count", "kind": "property", "signature": "GroupedGRUPopulationRunner.resident_parameter_count: int", "parameters": [], "line": 2112, "end_line": 2116, "return_type": "int", "decorators": [ "property" ] }, { "name": "patch_position", "kind": "property", "signature": "GroupedGRUPopulationRunner.patch_position: int", "parameters": [], "line": 2119, "end_line": 2120, "return_type": "int", "decorators": [ "property" ] }, { "name": "step", "kind": "method", "signature": "step(self, event_indices: torch.Tensor, event_values: torch.Tensor, morphology_indices: torch.Tensor) -> torch.Tensor", "parameters": [ { "name": "event_indices", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "event_values", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_indices", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 2122, "end_line": 2191, "return_type": "torch.Tensor", "decorators": [] }, { "name": "load_group_recurrent_parameters", "kind": "method", "signature": "load_group_recurrent_parameters(self, group: int, source: AxoTemporalModel) -> None", "parameters": [ { "name": "group", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "source", "type": "AxoTemporalModel", "default": null, "kind": "positional_or_keyword" } ], "line": 2193, "end_line": 2229, "return_type": "None", "decorators": [] } ], "bases": [] }, { "module": "deployment", "name": "NeuronBehaviorAdaptation", "kind": "class", "page": "api/population", "line": 7, "end_line": 135, "signature": "NeuronBehaviorAdaptation(width: int, branches: int, outputs: int, matrix_rank: int, adapt_recurrent: bool, branch_token_offset: bool = False) -> None", "parameters": [ { "name": "width", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "branches", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "outputs", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "matrix_rank", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "adapt_recurrent", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "branch_token_offset", "type": "bool", "default": "False", "kind": "positional_or_keyword" } ], "fields": [ { "name": "width", "type": "int", "default": null }, { "name": "branches", "type": "int", "default": null }, { "name": "outputs", "type": "int", "default": null }, { "name": "matrix_rank", "type": "int", "default": null }, { "name": "adapt_recurrent", "type": "bool", "default": null }, { "name": "branch_token_offset", "type": "bool", "default": "False" } ], "methods": [ { "name": "tensor_shapes", "kind": "property", "signature": "NeuronBehaviorAdaptation.tensor_shapes: dict[str, tuple[int, ...]]", "parameters": [], "line": 18, "end_line": 77, "return_type": "dict[str, tuple[int, ...]]", "decorators": [ "property" ] }, { "name": "tensor_slices", "kind": "property", "signature": "NeuronBehaviorAdaptation.tensor_slices: dict[str, slice]", "parameters": [], "line": 80, "end_line": 89, "return_type": "dict[str, slice]", "decorators": [ "property" ] }, { "name": "parameter_counts", "kind": "property", "signature": "NeuronBehaviorAdaptation.parameter_counts: dict[str, int]", "parameters": [], "line": 92, "end_line": 116, "return_type": "dict[str, int]", "decorators": [ "property" ] }, { "name": "parameter_count", "kind": "property", "signature": "NeuronBehaviorAdaptation.parameter_count: int", "parameters": [], "line": 119, "end_line": 120, "return_type": "int", "decorators": [ "property" ] }, { "name": "adapted_components", "kind": "property", "signature": "NeuronBehaviorAdaptation.adapted_components: frozenset[str]", "parameters": [], "line": 123, "end_line": 124, "return_type": "frozenset[str]", "decorators": [ "property" ] }, { "name": "shared_components", "kind": "property", "signature": "NeuronBehaviorAdaptation.shared_components: frozenset[str]", "parameters": [], "line": 127, "end_line": 135, "return_type": "frozenset[str]", "decorators": [ "property" ] } ], "bases": [], "inherited_fields": [] }, { "module": "deployment", "name": "PopulationDeploymentMeasurement", "kind": "class", "page": "api/population", "line": 139, "end_line": 165, "signature": "PopulationDeploymentMeasurement(population: int, adapted_parameters_per_neuron: int, adapted_components: tuple[str, ...], shared_components: tuple[str, ...], unique_adaptation_per_neuron: bool, adaptation_quality_validated: bool, shared_frozen_base: bool, output_materialized: bool, routing_feeds_model: bool, model_latency_ms: float, routing_latency_ms: float | None, complete_stack_latency_ms: float | None, model_peak_bytes: int, routing_peak_bytes: int | None, complete_stack_peak_bytes: int | None, local_source_identity_retained: bool = False, long_range_source_identity_retained: bool = False, morphology_assignment_independent: bool = False, activity_generation_mode: str = 'unknown', activity_dynamics_validated: bool = False, adaptation_execution_mode: str = 'direct', adaptation_reference_envelope_validated: bool = False, forecast_block_size: int = 1, complete_block_latency_ms: float | None = None) -> None", "parameters": [ { "name": "population", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "adapted_parameters_per_neuron", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "adapted_components", "type": "tuple[str, ...]", "default": null, "kind": "positional_or_keyword" }, { "name": "shared_components", "type": "tuple[str, ...]", "default": null, "kind": "positional_or_keyword" }, { "name": "unique_adaptation_per_neuron", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "adaptation_quality_validated", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "shared_frozen_base", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "output_materialized", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "routing_feeds_model", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "model_latency_ms", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "routing_latency_ms", "type": "float | None", "default": null, "kind": "positional_or_keyword" }, { "name": "complete_stack_latency_ms", "type": "float | None", "default": null, "kind": "positional_or_keyword" }, { "name": "model_peak_bytes", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "routing_peak_bytes", "type": "int | None", "default": null, "kind": "positional_or_keyword" }, { "name": "complete_stack_peak_bytes", "type": "int | None", "default": null, "kind": "positional_or_keyword" }, { "name": "local_source_identity_retained", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "long_range_source_identity_retained", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "morphology_assignment_independent", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "activity_generation_mode", "type": "str", "default": "'unknown'", "kind": "positional_or_keyword" }, { "name": "activity_dynamics_validated", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "adaptation_execution_mode", "type": "str", "default": "'direct'", "kind": "positional_or_keyword" }, { "name": "adaptation_reference_envelope_validated", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "forecast_block_size", "type": "int", "default": "1", "kind": "positional_or_keyword" }, { "name": "complete_block_latency_ms", "type": "float | None", "default": "None", "kind": "positional_or_keyword" } ], "fields": [ { "name": "population", "type": "int", "default": null }, { "name": "adapted_parameters_per_neuron", "type": "int", "default": null }, { "name": "adapted_components", "type": "tuple[str, ...]", "default": null }, { "name": "shared_components", "type": "tuple[str, ...]", "default": null }, { "name": "unique_adaptation_per_neuron", "type": "bool", "default": null }, { "name": "adaptation_quality_validated", "type": "bool", "default": null }, { "name": "shared_frozen_base", "type": "bool", "default": null }, { "name": "output_materialized", "type": "bool", "default": null }, { "name": "routing_feeds_model", "type": "bool", "default": null }, { "name": "model_latency_ms", "type": "float", "default": null }, { "name": "routing_latency_ms", "type": "float | None", "default": null }, { "name": "complete_stack_latency_ms", "type": "float | None", "default": null }, { "name": "model_peak_bytes", "type": "int", "default": null }, { "name": "routing_peak_bytes", "type": "int | None", "default": null }, { "name": "complete_stack_peak_bytes", "type": "int | None", "default": null }, { "name": "local_source_identity_retained", "type": "bool", "default": "False" }, { "name": "long_range_source_identity_retained", "type": "bool", "default": "False" }, { "name": "morphology_assignment_independent", "type": "bool", "default": "False" }, { "name": "activity_generation_mode", "type": "str", "default": "'unknown'" }, { "name": "activity_dynamics_validated", "type": "bool", "default": "False" }, { "name": "adaptation_execution_mode", "type": "str", "default": "'direct'" }, { "name": "adaptation_reference_envelope_validated", "type": "bool", "default": "False" }, { "name": "forecast_block_size", "type": "int", "default": "1" }, { "name": "complete_block_latency_ms", "type": "float | None", "default": "None" } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "deployment", "name": "PopulationDeploymentTarget", "kind": "class", "page": "api/population", "line": 169, "end_line": 376, "signature": "PopulationDeploymentTarget(headline_population: int, benchmark_population: int, biological_step_ms: float, adaptation: NeuronBehaviorAdaptation, shared_frozen_base: bool, output_materialization_required: bool, network_routing_required: bool, local_source_identity_required: bool, long_range_source_identity_required: bool, independent_morphology_assignment_required: bool, activity_dynamics_validation_required: bool, quota_activity_allowed: bool, forecast_block_size: int, firing_rate: float, fanout: int) -> None", "parameters": [ { "name": "headline_population", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "benchmark_population", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "biological_step_ms", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "adaptation", "type": "NeuronBehaviorAdaptation", "default": null, "kind": "positional_or_keyword" }, { "name": "shared_frozen_base", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "output_materialization_required", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "network_routing_required", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "local_source_identity_required", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "long_range_source_identity_required", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "independent_morphology_assignment_required", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "activity_dynamics_validation_required", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "quota_activity_allowed", "type": "bool", "default": null, "kind": "positional_or_keyword" }, { "name": "forecast_block_size", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "firing_rate", "type": "float", "default": null, "kind": "positional_or_keyword" }, { "name": "fanout", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "headline_population", "type": "int", "default": null }, { "name": "benchmark_population", "type": "int", "default": null }, { "name": "biological_step_ms", "type": "float", "default": null }, { "name": "adaptation", "type": "NeuronBehaviorAdaptation", "default": null }, { "name": "shared_frozen_base", "type": "bool", "default": null }, { "name": "output_materialization_required", "type": "bool", "default": null }, { "name": "network_routing_required", "type": "bool", "default": null }, { "name": "local_source_identity_required", "type": "bool", "default": null }, { "name": "long_range_source_identity_required", "type": "bool", "default": null }, { "name": "independent_morphology_assignment_required", "type": "bool", "default": null }, { "name": "activity_dynamics_validation_required", "type": "bool", "default": null }, { "name": "quota_activity_allowed", "type": "bool", "default": null }, { "name": "forecast_block_size", "type": "int", "default": null }, { "name": "firing_rate", "type": "float", "default": null }, { "name": "fanout", "type": "int", "default": null } ], "methods": [ { "name": "adapted_parameters_per_neuron", "kind": "property", "signature": "PopulationDeploymentTarget.adapted_parameters_per_neuron: int", "parameters": [], "line": 189, "end_line": 190, "return_type": "int", "decorators": [ "property" ] }, { "name": "validation_errors", "kind": "method", "signature": "validation_errors(self, measurement: PopulationDeploymentMeasurement) -> tuple[str, ...]", "parameters": [ { "name": "measurement", "type": "PopulationDeploymentMeasurement", "default": null, "kind": "positional_or_keyword" } ], "line": 192, "end_line": 340, "return_type": "tuple[str, ...]", "decorators": [] }, { "name": "evaluate", "kind": "method", "signature": "evaluate(self, measurement: PopulationDeploymentMeasurement) -> dict[str, int | float | bool]", "parameters": [ { "name": "measurement", "type": "PopulationDeploymentMeasurement", "default": null, "kind": "positional_or_keyword" } ], "line": 342, "end_line": 376, "return_type": "dict[str, int | float | bool]", "decorators": [] } ], "bases": [], "inherited_fields": [] }, { "module": "synapse", "name": "QuantizedSynapticEfficacyBank", "kind": "class", "page": "api/synapse", "line": 11, "end_line": 33, "signature": "QuantizedSynapticEfficacyBank(values: torch.Tensor, scale: torch.Tensor, channels_per_role: int, recurrent_stride: int, baseline: float = 1.0) -> None", "parameters": [ { "name": "values", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "scale", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "channels_per_role", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "recurrent_stride", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "baseline", "type": "float", "default": "1.0", "kind": "positional_or_keyword" } ], "fields": [ { "name": "values", "type": "torch.Tensor", "default": null }, { "name": "scale", "type": "torch.Tensor", "default": null }, { "name": "channels_per_role", "type": "int", "default": null }, { "name": "recurrent_stride", "type": "int", "default": null }, { "name": "baseline", "type": "float", "default": "1.0" } ], "methods": [ { "name": "population_size", "kind": "property", "signature": "QuantizedSynapticEfficacyBank.population_size: int", "parameters": [], "line": 21, "end_line": 22, "return_type": "int", "decorators": [ "property" ] }, { "name": "values_per_neuron", "kind": "property", "signature": "QuantizedSynapticEfficacyBank.values_per_neuron: int", "parameters": [], "line": 25, "end_line": 26, "return_type": "int", "decorators": [ "property" ] }, { "name": "storage_bytes", "kind": "property", "signature": "QuantizedSynapticEfficacyBank.storage_bytes: int", "parameters": [], "line": 29, "end_line": 33, "return_type": "int", "decorators": [ "property" ] } ], "bases": [], "inherited_fields": [] }, { "module": "synapse", "name": "retained_synaptic_efficacies_per_neuron", "kind": "function", "page": "api/synapse", "line": 36, "end_line": 50, "signature": "retained_synaptic_efficacies_per_neuron(*, channels_per_role: int, recurrent_stride: int) -> int", "parameters": [ { "name": "channels_per_role", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "recurrent_stride", "type": "int", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [] }, { "module": "synapse", "name": "quantize_synaptic_efficacies", "kind": "function", "page": "api/synapse", "line": 53, "end_line": 97, "signature": "quantize_synaptic_efficacies(efficacies: torch.Tensor, *, channels_per_role: int, recurrent_stride: int, baseline: float=1.0) -> QuantizedSynapticEfficacyBank", "parameters": [ { "name": "efficacies", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "channels_per_role", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "recurrent_stride", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "baseline", "type": "float", "default": "1.0", "kind": "keyword_only", "resolved_default": 1.0 } ], "fields": [], "methods": [] }, { "module": "synapse", "name": "dequantize_synaptic_efficacies", "kind": "function", "page": "api/synapse", "line": 100, "end_line": 109, "signature": "dequantize_synaptic_efficacies(bank: QuantizedSynapticEfficacyBank) -> torch.Tensor", "parameters": [ { "name": "bank", "type": "QuantizedSynapticEfficacyBank", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "connectome", "name": "DelayBucket", "kind": "class", "page": "api/connectome", "line": 24, "end_line": 35, "signature": "DelayBucket(delay_steps: int, first_fanout_offset: int, event_count_per_source: int, local_targets: bool) -> None", "parameters": [ { "name": "delay_steps", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "first_fanout_offset", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "event_count_per_source", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "local_targets", "type": "bool", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "delay_steps", "type": "int", "default": null }, { "name": "first_fanout_offset", "type": "int", "default": null }, { "name": "event_count_per_source", "type": "int", "default": null }, { "name": "local_targets", "type": "bool", "default": null } ], "methods": [ { "name": "fanout_offsets", "kind": "property", "signature": "DelayBucket.fanout_offsets: tuple[int, ...]", "parameters": [], "line": 31, "end_line": 35, "return_type": "tuple[int, ...]", "decorators": [ "property" ] } ], "bases": [], "inherited_fields": [] }, { "module": "connectome", "name": "procedural_delay_buckets", "kind": "function", "page": "api/connectome", "line": 38, "end_line": 71, "signature": "procedural_delay_buckets(*, fanout: int, local_fanout: int, minimum_delay_steps: int=1) -> tuple[DelayBucket, ...]", "parameters": [ { "name": "fanout", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "local_fanout", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "minimum_delay_steps", "type": "int", "default": "1", "kind": "keyword_only", "resolved_default": 1 } ], "fields": [], "methods": [] }, { "module": "connectome", "name": "compact_active_sources", "kind": "function", "page": "api/connectome", "line": 74, "end_line": 86, "signature": "compact_active_sources(outputs: torch.Tensor, *, threshold: float=0.0) -> torch.Tensor", "parameters": [ { "name": "outputs", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "threshold", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 } ], "fields": [], "methods": [] }, { "module": "connectome", "name": "build_typed_tile_pools", "kind": "function", "page": "api/connectome", "line": 89, "end_line": 164, "signature": "build_typed_tile_pools(source_is_inhibitory: torch.Tensor, physical_to_logical: torch.Tensor, *, spatial_tile_neurons: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]", "parameters": [ { "name": "source_is_inhibitory", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "physical_to_logical", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "spatial_tile_neurons", "type": "int", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [] }, { "module": "connectome", "name": "build_external_count_branch_bank", "kind": "function", "page": "api/connectome", "line": 167, "end_line": 261, "signature": "build_external_count_branch_bank(*, slot_map: torch.Tensor, morphology_gain: torch.Tensor, feature_gain: torch.Tensor, synapses_per_branch: int, max_event_count: int=32, channels_per_type: int | None=None) -> torch.Tensor", "parameters": [ { "name": "slot_map", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "morphology_gain", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "feature_gain", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "synapses_per_branch", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "max_event_count", "type": "int", "default": "32", "kind": "keyword_only", "resolved_default": 32 }, { "name": "channels_per_type", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "fields": [], "methods": [] }, { "module": "connectome", "name": "TrajectoryExternalBranchDrive", "kind": "class", "page": "api/connectome", "line": 1514, "end_line": 1768, "signature": "TrajectoryExternalBranchDrive(self, *, trajectories: torch.Tensor, branch_bank: torch.Tensor, logical_ids: torch.Tensor, morphology: torch.Tensor, seed: int, excitatory_gain: float, inhibitory_gain: float, slot_map: torch.Tensor | None=None, morphology_gain: torch.Tensor | None=None, feature_gain: torch.Tensor | None=None, synapses_per_branch: int | None=None, block_rows: int=8, block_branches: int=128) -> None", "parameters": [ { "name": "trajectories", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "branch_bank", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "logical_ids", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "morphology", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "seed", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "excitatory_gain", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "inhibitory_gain", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "slot_map", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "morphology_gain", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "feature_gain", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "synapses_per_branch", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "block_rows", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "block_branches", "type": "int", "default": "128", "kind": "keyword_only", "resolved_default": 128 } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, *, trajectories: torch.Tensor, branch_bank: torch.Tensor, logical_ids: torch.Tensor, morphology: torch.Tensor, seed: int, excitatory_gain: float, inhibitory_gain: float, slot_map: torch.Tensor | None=None, morphology_gain: torch.Tensor | None=None, feature_gain: torch.Tensor | None=None, synapses_per_branch: int | None=None, block_rows: int=8, block_branches: int=128) -> None", "parameters": [ { "name": "trajectories", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "branch_bank", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "logical_ids", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "morphology", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "seed", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "excitatory_gain", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "inhibitory_gain", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "slot_map", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "morphology_gain", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "feature_gain", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "synapses_per_branch", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "block_rows", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "block_branches", "type": "int", "default": "128", "kind": "keyword_only", "resolved_default": 128 } ], "line": 1517, "end_line": 1640, "return_type": "None", "decorators": [] }, { "name": "materialize", "kind": "method", "signature": "materialize(self, step: int, *, output: torch.Tensor | None=None, add_existing: bool=False, existing: torch.Tensor | None=None, recurrent_bits: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "step", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "output", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "add_existing", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "existing", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "recurrent_bits", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 1642, "end_line": 1768, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [] }, { "module": "connectome", "name": "ProceduralBinaryChannelConnectomeRouter", "kind": "class", "page": "api/connectome", "line": 1771, "end_line": 2083, "signature": "ProceduralBinaryChannelConnectomeRouter(self, contract: LargePopulationSimulationContract, *, slot_map: torch.Tensor, morphology: torch.Tensor, source_is_inhibitory: torch.Tensor, morphology_gain: torch.Tensor, feature_gain: torch.Tensor, synapses_per_branch: int, physical_to_logical: torch.Tensor | None=None, logical_to_physical: torch.Tensor | None=None, route_block_size: int=512, route_num_warps: int=4, branch_block_rows: int=8, branch_block_size: int=128) -> None", "parameters": [ { "name": "contract", "type": "LargePopulationSimulationContract", "default": null, "kind": "positional_or_keyword" }, { "name": "slot_map", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "morphology", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "source_is_inhibitory", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "morphology_gain", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "feature_gain", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "synapses_per_branch", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "physical_to_logical", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "logical_to_physical", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "route_block_size", "type": "int", "default": "512", "kind": "keyword_only", "resolved_default": 512 }, { "name": "route_num_warps", "type": "int", "default": "4", "kind": "keyword_only", "resolved_default": 4 }, { "name": "branch_block_rows", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "branch_block_size", "type": "int", "default": "128", "kind": "keyword_only", "resolved_default": 128 } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, contract: LargePopulationSimulationContract, *, slot_map: torch.Tensor, morphology: torch.Tensor, source_is_inhibitory: torch.Tensor, morphology_gain: torch.Tensor, feature_gain: torch.Tensor, synapses_per_branch: int, physical_to_logical: torch.Tensor | None=None, logical_to_physical: torch.Tensor | None=None, route_block_size: int=512, route_num_warps: int=4, branch_block_rows: int=8, branch_block_size: int=128) -> None", "parameters": [ { "name": "contract", "type": "LargePopulationSimulationContract", "default": null, "kind": "positional_or_keyword" }, { "name": "slot_map", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "morphology", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "source_is_inhibitory", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "morphology_gain", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "feature_gain", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "synapses_per_branch", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "physical_to_logical", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "logical_to_physical", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "route_block_size", "type": "int", "default": "512", "kind": "keyword_only", "resolved_default": 512 }, { "name": "route_num_warps", "type": "int", "default": "4", "kind": "keyword_only", "resolved_default": 4 }, { "name": "branch_block_rows", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "branch_block_size", "type": "int", "default": "128", "kind": "keyword_only", "resolved_default": 128 } ], "line": 1774, "end_line": 1949, "return_type": "None", "decorators": [] }, { "name": "queue_bytes", "kind": "property", "signature": "ProceduralBinaryChannelConnectomeRouter.queue_bytes: int", "parameters": [], "line": 1952, "end_line": 1953, "return_type": "int", "decorators": [ "property" ] }, { "name": "current_channel_bits", "kind": "method", "signature": "current_channel_bits(self, current_slot: int) -> torch.Tensor", "parameters": [ { "name": "current_slot", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "line": 1955, "end_line": 1957, "return_type": "torch.Tensor", "decorators": [] }, { "name": "current_inputs", "kind": "method", "signature": "current_inputs(self, current_slot: int, *, output: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "current_slot", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "output", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 1959, "end_line": 2018, "return_type": "torch.Tensor", "decorators": [] }, { "name": "clear_and_route", "kind": "method", "signature": "clear_and_route(self, active_sources: torch.Tensor, *, current_slot: int, clear_consumed: bool=True) -> torch.Tensor", "parameters": [ { "name": "active_sources", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "current_slot", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "clear_consumed", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "line": 2020, "end_line": 2079, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [] }, { "module": "connectome", "name": "ProceduralExactBranchConnectomeRouter", "kind": "class", "page": "api/connectome", "line": 2086, "end_line": 2277, "signature": "ProceduralExactBranchConnectomeRouter(self, *args, external_trajectories: torch.Tensor | None=None, external_seed: int=0, external_excitatory_gain: float=1.0, external_inhibitory_gain: float=1.0, max_external_event_count: int=32, **kwargs) -> None", "parameters": [ { "name": "args", "type": null, "default": null, "kind": "variadic_positional" }, { "name": "external_trajectories", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "external_seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "external_excitatory_gain", "type": "float", "default": "1.0", "kind": "keyword_only", "resolved_default": 1.0 }, { "name": "external_inhibitory_gain", "type": "float", "default": "1.0", "kind": "keyword_only", "resolved_default": 1.0 }, { "name": "max_external_event_count", "type": "int", "default": "32", "kind": "keyword_only", "resolved_default": 32 }, { "name": "kwargs", "type": null, "default": null, "kind": "variadic_keyword" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, *args, external_trajectories: torch.Tensor | None=None, external_seed: int=0, external_excitatory_gain: float=1.0, external_inhibitory_gain: float=1.0, max_external_event_count: int=32, **kwargs) -> None", "parameters": [ { "name": "args", "type": null, "default": null, "kind": "variadic_positional" }, { "name": "external_trajectories", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "external_seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "external_excitatory_gain", "type": "float", "default": "1.0", "kind": "keyword_only", "resolved_default": 1.0 }, { "name": "external_inhibitory_gain", "type": "float", "default": "1.0", "kind": "keyword_only", "resolved_default": 1.0 }, { "name": "max_external_event_count", "type": "int", "default": "32", "kind": "keyword_only", "resolved_default": 32 }, { "name": "kwargs", "type": null, "default": null, "kind": "variadic_keyword" } ], "line": 2091, "end_line": 2147, "return_type": "None", "decorators": [] }, { "name": "queue_bytes", "kind": "property", "signature": "ProceduralExactBranchConnectomeRouter.queue_bytes: int", "parameters": [], "line": 2150, "end_line": 2156, "return_type": "int", "decorators": [ "property" ] }, { "name": "current_inputs", "kind": "method", "signature": "current_inputs(self, current_slot: int, *, output: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "current_slot", "type": "int", "default": null, "kind": "positional_or_keyword" }, { "name": "output", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 2158, "end_line": 2179, "return_type": "torch.Tensor", "decorators": [] }, { "name": "clear_and_route", "kind": "method", "signature": "clear_and_route(self, active_sources: torch.Tensor, *, current_slot: int, current_step: int | None=None, clear_consumed: bool=True) -> torch.Tensor", "parameters": [ { "name": "active_sources", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "current_slot", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "current_step", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "clear_consumed", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "line": 2181, "end_line": 2277, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [ "ProceduralBinaryChannelConnectomeRouter" ] }, { "module": "connectome", "name": "ProceduralMorphologyConnectomeRouter", "kind": "class", "page": "api/connectome", "line": 2280, "end_line": 2606, "signature": "ProceduralMorphologyConnectomeRouter(self, contract: LargePopulationSimulationContract, *, slot_map: torch.Tensor, morphology: torch.Tensor, source_is_inhibitory: torch.Tensor, morphology_gain: torch.Tensor, feature_gain: torch.Tensor, synapses_per_branch: int, physical_to_logical: torch.Tensor | None=None, logical_to_physical: torch.Tensor | None=None, route_block_size: int=512, route_num_warps: int=4) -> None", "parameters": [ { "name": "contract", "type": "LargePopulationSimulationContract", "default": null, "kind": "positional_or_keyword" }, { "name": "slot_map", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "morphology", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "source_is_inhibitory", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "morphology_gain", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "feature_gain", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "synapses_per_branch", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "physical_to_logical", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "logical_to_physical", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "route_block_size", "type": "int", "default": "512", "kind": "keyword_only", "resolved_default": 512 }, { "name": "route_num_warps", "type": "int", "default": "4", "kind": "keyword_only", "resolved_default": 4 } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, contract: LargePopulationSimulationContract, *, slot_map: torch.Tensor, morphology: torch.Tensor, source_is_inhibitory: torch.Tensor, morphology_gain: torch.Tensor, feature_gain: torch.Tensor, synapses_per_branch: int, physical_to_logical: torch.Tensor | None=None, logical_to_physical: torch.Tensor | None=None, route_block_size: int=512, route_num_warps: int=4) -> None", "parameters": [ { "name": "contract", "type": "LargePopulationSimulationContract", "default": null, "kind": "positional_or_keyword" }, { "name": "slot_map", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "morphology", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "source_is_inhibitory", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "morphology_gain", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "feature_gain", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "synapses_per_branch", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "physical_to_logical", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "logical_to_physical", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "route_block_size", "type": "int", "default": "512", "kind": "keyword_only", "resolved_default": 512 }, { "name": "route_num_warps", "type": "int", "default": "4", "kind": "keyword_only", "resolved_default": 4 } ], "line": 2283, "end_line": 2445, "return_type": "None", "decorators": [] }, { "name": "current_inputs", "kind": "method", "signature": "current_inputs(self, current_slot: int) -> torch.Tensor", "parameters": [ { "name": "current_slot", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "line": 2447, "end_line": 2449, "return_type": "torch.Tensor", "decorators": [] }, { "name": "clear_and_route", "kind": "method", "signature": "clear_and_route(self, active_sources: torch.Tensor, *, current_slot: int, clear_consumed: bool=True, touched_offsets: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "active_sources", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "current_slot", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "clear_consumed", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "touched_offsets", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 2451, "end_line": 2577, "return_type": "torch.Tensor", "decorators": [] }, { "name": "clear_recorded_branches", "kind": "method", "signature": "clear_recorded_branches(self, touched_offsets: torch.Tensor) -> None", "parameters": [ { "name": "touched_offsets", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 2579, "end_line": 2602, "return_type": "None", "decorators": [] } ], "bases": [] }, { "module": "connectome", "name": "ProceduralUniqueChannelConnectomeRouter", "kind": "class", "page": "api/connectome", "line": 2609, "end_line": 3007, "signature": "ProceduralUniqueChannelConnectomeRouter(self, *args, synaptic_efficacy_bank: QuantizedSynapticEfficacyBank | None=None, external_trajectories: torch.Tensor | None=None, external_seed: int=0, external_excitatory_gain: float=1.0, external_inhibitory_gain: float=1.0, max_external_event_count: int=32, allow_repeated_source_target_pairs: bool=False, **kwargs) -> None", "parameters": [ { "name": "args", "type": null, "default": null, "kind": "variadic_positional" }, { "name": "synaptic_efficacy_bank", "type": "QuantizedSynapticEfficacyBank | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "external_trajectories", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "external_seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "external_excitatory_gain", "type": "float", "default": "1.0", "kind": "keyword_only", "resolved_default": 1.0 }, { "name": "external_inhibitory_gain", "type": "float", "default": "1.0", "kind": "keyword_only", "resolved_default": 1.0 }, { "name": "max_external_event_count", "type": "int", "default": "32", "kind": "keyword_only", "resolved_default": 32 }, { "name": "allow_repeated_source_target_pairs", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "kwargs", "type": null, "default": null, "kind": "variadic_keyword" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, *args, synaptic_efficacy_bank: QuantizedSynapticEfficacyBank | None=None, external_trajectories: torch.Tensor | None=None, external_seed: int=0, external_excitatory_gain: float=1.0, external_inhibitory_gain: float=1.0, max_external_event_count: int=32, allow_repeated_source_target_pairs: bool=False, **kwargs) -> None", "parameters": [ { "name": "args", "type": null, "default": null, "kind": "variadic_positional" }, { "name": "synaptic_efficacy_bank", "type": "QuantizedSynapticEfficacyBank | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "external_trajectories", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "external_seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "external_excitatory_gain", "type": "float", "default": "1.0", "kind": "keyword_only", "resolved_default": 1.0 }, { "name": "external_inhibitory_gain", "type": "float", "default": "1.0", "kind": "keyword_only", "resolved_default": 1.0 }, { "name": "max_external_event_count", "type": "int", "default": "32", "kind": "keyword_only", "resolved_default": 32 }, { "name": "allow_repeated_source_target_pairs", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "kwargs", "type": null, "default": null, "kind": "variadic_keyword" } ], "line": 2614, "end_line": 2753, "return_type": "None", "decorators": [] }, { "name": "queue_bytes", "kind": "property", "signature": "ProceduralUniqueChannelConnectomeRouter.queue_bytes: int", "parameters": [], "line": 2756, "end_line": 2757, "return_type": "int", "decorators": [ "property" ] }, { "name": "recorded_offsets_per_source", "kind": "property", "signature": "ProceduralUniqueChannelConnectomeRouter.recorded_offsets_per_source: int", "parameters": [], "line": 2760, "end_line": 2779, "return_type": "int", "decorators": [ "property" ] }, { "name": "clear_and_route", "kind": "method", "signature": "clear_and_route(self, active_sources: torch.Tensor, *, current_slot: int, current_step: int | None=None, clear_consumed: bool=True, touched_offsets: torch.Tensor | None=None, unique_rows: torch.Tensor | None=None, unique_row_flags: torch.Tensor | None=None, unique_row_count: torch.Tensor | None=None) -> torch.Tensor", "parameters": [ { "name": "active_sources", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "current_slot", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "current_step", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "clear_consumed", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true }, { "name": "touched_offsets", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "unique_rows", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "unique_row_flags", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "unique_row_count", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "line": 2781, "end_line": 3007, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [ "ProceduralMorphologyConnectomeRouter" ] }, { "module": "activity", "name": "StratifiedPatchPhaseLayout", "kind": "class", "page": "api/activity", "line": 12, "end_line": 32, "signature": "StratifiedPatchPhaseLayout(phase_ids: torch.Tensor, storage_permutation: torch.Tensor, inverse_permutation: torch.Tensor, phase_population_sizes: tuple[int, ...], storage_group_offsets: tuple[int, ...], storage_group_shape: tuple[int, int]) -> None", "parameters": [ { "name": "phase_ids", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "storage_permutation", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "inverse_permutation", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "phase_population_sizes", "type": "tuple[int, ...]", "default": null, "kind": "positional_or_keyword" }, { "name": "storage_group_offsets", "type": "tuple[int, ...]", "default": null, "kind": "positional_or_keyword" }, { "name": "storage_group_shape", "type": "tuple[int, int]", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "phase_ids", "type": "torch.Tensor", "default": null }, { "name": "storage_permutation", "type": "torch.Tensor", "default": null }, { "name": "inverse_permutation", "type": "torch.Tensor", "default": null }, { "name": "phase_population_sizes", "type": "tuple[int, ...]", "default": null }, { "name": "storage_group_offsets", "type": "tuple[int, ...]", "default": null }, { "name": "storage_group_shape", "type": "tuple[int, int]", "default": null } ], "methods": [ { "name": "patch_size", "kind": "property", "signature": "StratifiedPatchPhaseLayout.patch_size: int", "parameters": [], "line": 23, "end_line": 24, "return_type": "int", "decorators": [ "property" ] }, { "name": "population_size", "kind": "property", "signature": "StratifiedPatchPhaseLayout.population_size: int", "parameters": [], "line": 27, "end_line": 28, "return_type": "int", "decorators": [ "property" ] }, { "name": "phase_sizes", "kind": "property", "signature": "StratifiedPatchPhaseLayout.phase_sizes: tuple[int, ...]", "parameters": [], "line": 31, "end_line": 32, "return_type": "tuple[int, ...]", "decorators": [ "property" ] } ], "bases": [], "inherited_fields": [] }, { "module": "activity", "name": "FixedActivityThresholds", "kind": "class", "page": "api/activity", "line": 36, "end_line": 67, "signature": "FixedActivityThresholds(excitatory: tuple[float, ...], inhibitory: tuple[float, ...]) -> None", "parameters": [ { "name": "excitatory", "type": "tuple[float, ...]", "default": null, "kind": "positional_or_keyword" }, { "name": "inhibitory", "type": "tuple[float, ...]", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "excitatory", "type": "tuple[float, ...]", "default": null }, { "name": "inhibitory", "type": "tuple[float, ...]", "default": null } ], "methods": [ { "name": "morphology_count", "kind": "property", "signature": "FixedActivityThresholds.morphology_count: int", "parameters": [], "line": 54, "end_line": 55, "return_type": "int", "decorators": [ "property" ] }, { "name": "as_tensor", "kind": "method", "signature": "as_tensor(self, *, device: torch.device | str, dtype: torch.dtype) -> torch.Tensor", "parameters": [ { "name": "device", "type": "torch.device | str", "default": null, "kind": "keyword_only" }, { "name": "dtype", "type": "torch.dtype", "default": null, "kind": "keyword_only" } ], "line": 57, "end_line": 67, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [], "inherited_fields": [] }, { "module": "activity", "name": "FixedBlockActivityThresholds", "kind": "class", "page": "api/activity", "line": 71, "end_line": 106, "signature": "FixedBlockActivityThresholds(forecast_steps: tuple[FixedActivityThresholds, ...]) -> None", "parameters": [ { "name": "forecast_steps", "type": "tuple[FixedActivityThresholds, ...]", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "forecast_steps", "type": "tuple[FixedActivityThresholds, ...]", "default": null } ], "methods": [ { "name": "forecast_step_count", "kind": "property", "signature": "FixedBlockActivityThresholds.forecast_step_count: int", "parameters": [], "line": 88, "end_line": 89, "return_type": "int", "decorators": [ "property" ] }, { "name": "morphology_count", "kind": "property", "signature": "FixedBlockActivityThresholds.morphology_count: int", "parameters": [], "line": 92, "end_line": 93, "return_type": "int", "decorators": [ "property" ] }, { "name": "as_tensor", "kind": "method", "signature": "as_tensor(self, *, device: torch.device | str, dtype: torch.dtype) -> torch.Tensor", "parameters": [ { "name": "device", "type": "torch.device | str", "default": null, "kind": "keyword_only" }, { "name": "dtype", "type": "torch.dtype", "default": null, "kind": "keyword_only" } ], "line": 95, "end_line": 106, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [], "inherited_fields": [] }, { "module": "activity", "name": "HomeostaticThresholdController", "kind": "class", "page": "api/activity", "line": 109, "end_line": 253, "signature": "HomeostaticThresholdController(self, *, group_ids: torch.Tensor, target_rates_hz: torch.Tensor, dt_ms: float, update_interval_ms: float, time_constant_ms: float, learning_rate: float, max_abs_offset: float) -> None", "parameters": [ { "name": "group_ids", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "target_rates_hz", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "dt_ms", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "update_interval_ms", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "time_constant_ms", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "learning_rate", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "max_abs_offset", "type": "float", "default": null, "kind": "keyword_only" } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, *, group_ids: torch.Tensor, target_rates_hz: torch.Tensor, dt_ms: float, update_interval_ms: float, time_constant_ms: float, learning_rate: float, max_abs_offset: float) -> None", "parameters": [ { "name": "group_ids", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "target_rates_hz", "type": "torch.Tensor", "default": null, "kind": "keyword_only" }, { "name": "dt_ms", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "update_interval_ms", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "time_constant_ms", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "learning_rate", "type": "float", "default": null, "kind": "keyword_only" }, { "name": "max_abs_offset", "type": "float", "default": null, "kind": "keyword_only" } ], "line": 117, "end_line": 204, "return_type": "None", "decorators": [] }, { "name": "observe", "kind": "method", "signature": "observe(self, activity: torch.Tensor) -> bool", "parameters": [ { "name": "activity", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" } ], "line": 207, "end_line": 248, "return_type": "bool", "decorators": [ "torch.no_grad()" ] }, { "name": "threshold_offsets_per_neuron", "kind": "method", "signature": "threshold_offsets_per_neuron(self) -> torch.Tensor", "parameters": [], "line": 250, "end_line": 253, "return_type": "torch.Tensor", "decorators": [] } ], "bases": [] }, { "module": "activity", "name": "build_stratified_patch_phase_layout", "kind": "function", "page": "api/activity", "line": 256, "end_line": 376, "signature": "build_stratified_patch_phase_layout(morphology_ids: torch.Tensor, inhibitory: torch.Tensor, *, patch_size: int, seed: int, tile_ids: torch.Tensor | None=None) -> StratifiedPatchPhaseLayout", "parameters": [ { "name": "morphology_ids", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "inhibitory", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "patch_size", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "seed", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "tile_ids", "type": "torch.Tensor | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "fields": [], "methods": [] }, { "module": "activity", "name": "threshold_model_activity", "kind": "function", "page": "api/activity", "line": 379, "end_line": 414, "signature": "threshold_model_activity(spike_scores: torch.Tensor, morphology_ids: torch.Tensor, inhibitory: torch.Tensor, thresholds: FixedActivityThresholds) -> torch.Tensor", "parameters": [ { "name": "spike_scores", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_ids", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "inhibitory", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "thresholds", "type": "FixedActivityThresholds", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "activity", "name": "threshold_block_model_activity", "kind": "function", "page": "api/activity", "line": 417, "end_line": 451, "signature": "threshold_block_model_activity(spike_scores: torch.Tensor, morphology_ids: torch.Tensor, inhibitory: torch.Tensor, thresholds: FixedBlockActivityThresholds) -> torch.Tensor", "parameters": [ { "name": "spike_scores", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "morphology_ids", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "inhibitory", "type": "torch.Tensor", "default": null, "kind": "positional_or_keyword" }, { "name": "thresholds", "type": "FixedBlockActivityThresholds", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "data", "name": "NeuronIOSample", "kind": "class", "page": "api/data", "line": 14, "end_line": 17, "signature": "NeuronIOSample(sample_id: str, inputs: np.ndarray, targets: np.ndarray) -> None", "parameters": [ { "name": "sample_id", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "inputs", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" }, { "name": "targets", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "sample_id", "type": "str", "default": null }, { "name": "inputs", "type": "np.ndarray", "default": null }, { "name": "targets", "type": "np.ndarray", "default": null } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "data", "name": "ShardedNeuronIODataset", "kind": "class", "page": "api/data", "line": 29, "end_line": 222, "signature": "ShardedNeuronIODataset(self, root: str | Path, *, shuffle: bool=False, seed: int=0, shuffle_mode: Literal['sample', 'shard']='sample', cache_shards: int=0) -> None", "parameters": [ { "name": "root", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "shuffle", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "shuffle_mode", "type": "Literal['sample', 'shard']", "default": "'sample'", "kind": "keyword_only", "resolved_default": "sample" }, { "name": "cache_shards", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, root: str | Path, *, shuffle: bool=False, seed: int=0, shuffle_mode: Literal['sample', 'shard']='sample', cache_shards: int=0) -> None", "parameters": [ { "name": "root", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "shuffle", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "shuffle_mode", "type": "Literal['sample', 'shard']", "default": "'sample'", "kind": "keyword_only", "resolved_default": "sample" }, { "name": "cache_shards", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 } ], "line": 32, "end_line": 67, "return_type": "None", "decorators": [] }, { "name": "__len__", "kind": "method", "signature": "__len__(self) -> int", "parameters": [], "line": 69, "end_line": 70, "return_type": "int", "decorators": [] }, { "name": "reshuffle", "kind": "method", "signature": "reshuffle(self, seed: int) -> None", "parameters": [ { "name": "seed", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "line": 72, "end_line": 80, "return_type": "None", "decorators": [] }, { "name": "__getitem__", "kind": "method", "signature": "__getitem__(self, index: int) -> NeuronIOSample", "parameters": [ { "name": "index", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "line": 82, "end_line": 91, "return_type": "NeuronIOSample", "decorators": [] }, { "name": "get_batch", "kind": "method", "signature": "get_batch(self, indices: range | list[int]) -> tuple[np.ndarray, np.ndarray]", "parameters": [ { "name": "indices", "type": "range | list[int]", "default": null, "kind": "positional_or_keyword" } ], "line": 93, "end_line": 101, "return_type": "tuple[np.ndarray, np.ndarray]", "decorators": [] }, { "name": "get_sample_ids", "kind": "method", "signature": "get_sample_ids(self, indices: range | list[int]) -> list[str]", "parameters": [ { "name": "indices", "type": "range | list[int]", "default": null, "kind": "positional_or_keyword" } ], "line": 103, "end_line": 104, "return_type": "list[str]", "decorators": [] }, { "name": "get_optional_array_batch", "kind": "method", "signature": "get_optional_array_batch(self, name: str, indices: range | list[int]) -> np.ndarray | None", "parameters": [ { "name": "name", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "indices", "type": "range | list[int]", "default": null, "kind": "positional_or_keyword" } ], "line": 106, "end_line": 115, "return_type": "np.ndarray | None", "decorators": [] }, { "name": "get_shard_sequence_length", "kind": "method", "signature": "get_shard_sequence_length(self, shard_idx: int) -> int", "parameters": [ { "name": "shard_idx", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "line": 117, "end_line": 130, "return_type": "int", "decorators": [] } ], "bases": [] }, { "module": "data", "name": "DeterministicWindowDataset", "kind": "class", "page": "api/data", "line": 225, "end_line": 280, "signature": "DeterministicWindowDataset(self, base: ShardedNeuronIODataset, *, window_size: int, stride: int | None=None, start_offset: int=0) -> None", "parameters": [ { "name": "base", "type": "ShardedNeuronIODataset", "default": null, "kind": "positional_or_keyword" }, { "name": "window_size", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "stride", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "start_offset", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, base: ShardedNeuronIODataset, *, window_size: int, stride: int | None=None, start_offset: int=0) -> None", "parameters": [ { "name": "base", "type": "ShardedNeuronIODataset", "default": null, "kind": "positional_or_keyword" }, { "name": "window_size", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "stride", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "start_offset", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 } ], "line": 228, "end_line": 257, "return_type": "None", "decorators": [] }, { "name": "__len__", "kind": "method", "signature": "__len__(self) -> int", "parameters": [], "line": 259, "end_line": 260, "return_type": "int", "decorators": [] }, { "name": "__getitem__", "kind": "method", "signature": "__getitem__(self, index: int) -> NeuronIOSample", "parameters": [ { "name": "index", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "line": 262, "end_line": 270, "return_type": "NeuronIOSample", "decorators": [] }, { "name": "get_batch", "kind": "method", "signature": "get_batch(self, indices: range | list[int]) -> tuple[np.ndarray, np.ndarray]", "parameters": [ { "name": "indices", "type": "range | list[int]", "default": null, "kind": "positional_or_keyword" } ], "line": 272, "end_line": 277, "return_type": "tuple[np.ndarray, np.ndarray]", "decorators": [] }, { "name": "get_sample_ids", "kind": "method", "signature": "get_sample_ids(self, indices: range | list[int]) -> list[str]", "parameters": [ { "name": "indices", "type": "range | list[int]", "default": null, "kind": "positional_or_keyword" } ], "line": 279, "end_line": 280, "return_type": "list[str]", "decorators": [] } ], "bases": [] }, { "module": "data", "name": "RandomFullTraceWindowDataset", "kind": "class", "page": "api/data", "line": 283, "end_line": 399, "signature": "RandomFullTraceWindowDataset(self, base: ShardedNeuronIODataset, *, window_size: int=500, start_offset: int=500, samples_per_epoch: int | None=None, batch_size: int=8, shard_reuse_batches: int=1, sequence_length: int | None=None, seed: int=0, cache_full_shards: bool=True) -> None", "parameters": [ { "name": "base", "type": "ShardedNeuronIODataset", "default": null, "kind": "positional_or_keyword" }, { "name": "window_size", "type": "int", "default": "500", "kind": "keyword_only", "resolved_default": 500 }, { "name": "start_offset", "type": "int", "default": "500", "kind": "keyword_only", "resolved_default": 500 }, { "name": "samples_per_epoch", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "batch_size", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "shard_reuse_batches", "type": "int", "default": "1", "kind": "keyword_only", "resolved_default": 1 }, { "name": "sequence_length", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "cache_full_shards", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, base: ShardedNeuronIODataset, *, window_size: int=500, start_offset: int=500, samples_per_epoch: int | None=None, batch_size: int=8, shard_reuse_batches: int=1, sequence_length: int | None=None, seed: int=0, cache_full_shards: bool=True) -> None", "parameters": [ { "name": "base", "type": "ShardedNeuronIODataset", "default": null, "kind": "positional_or_keyword" }, { "name": "window_size", "type": "int", "default": "500", "kind": "keyword_only", "resolved_default": 500 }, { "name": "start_offset", "type": "int", "default": "500", "kind": "keyword_only", "resolved_default": 500 }, { "name": "samples_per_epoch", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "batch_size", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "shard_reuse_batches", "type": "int", "default": "1", "kind": "keyword_only", "resolved_default": 1 }, { "name": "sequence_length", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "cache_full_shards", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "line": 286, "end_line": 326, "return_type": "None", "decorators": [] }, { "name": "__len__", "kind": "method", "signature": "__len__(self) -> int", "parameters": [], "line": 328, "end_line": 329, "return_type": "int", "decorators": [] }, { "name": "reshuffle", "kind": "method", "signature": "reshuffle(self, seed: int) -> None", "parameters": [ { "name": "seed", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "line": 331, "end_line": 354, "return_type": "None", "decorators": [] }, { "name": "__getitem__", "kind": "method", "signature": "__getitem__(self, index: int) -> NeuronIOSample", "parameters": [ { "name": "index", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "line": 356, "end_line": 364, "return_type": "NeuronIOSample", "decorators": [] }, { "name": "get_batch", "kind": "method", "signature": "get_batch(self, indices: range | list[int]) -> tuple[np.ndarray, np.ndarray]", "parameters": [ { "name": "indices", "type": "range | list[int]", "default": null, "kind": "positional_or_keyword" } ], "line": 366, "end_line": 383, "return_type": "tuple[np.ndarray, np.ndarray]", "decorators": [] }, { "name": "get_sample_ids", "kind": "method", "signature": "get_sample_ids(self, indices: range | list[int]) -> list[str]", "parameters": [ { "name": "indices", "type": "range | list[int]", "default": null, "kind": "positional_or_keyword" } ], "line": 385, "end_line": 386, "return_type": "list[str]", "decorators": [] } ], "bases": [] }, { "module": "data", "name": "OfficialStyleFullTraceWindowDataset", "kind": "class", "page": "api/data", "line": 402, "end_line": 543, "signature": "OfficialStyleFullTraceWindowDataset(self, base: ShardedNeuronIODataset, *, window_size: int=500, start_offset: int=500, samples_per_epoch: int | None=None, batch_size: int=8, file_load_fraction: float=0.3, source_simulations: int=128, sequence_length: int | None=None, seed: int=0, cache_full_shards: bool=True) -> None", "parameters": [ { "name": "base", "type": "ShardedNeuronIODataset", "default": null, "kind": "positional_or_keyword" }, { "name": "window_size", "type": "int", "default": "500", "kind": "keyword_only", "resolved_default": 500 }, { "name": "start_offset", "type": "int", "default": "500", "kind": "keyword_only", "resolved_default": 500 }, { "name": "samples_per_epoch", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "batch_size", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "file_load_fraction", "type": "float", "default": "0.3", "kind": "keyword_only", "resolved_default": 0.3 }, { "name": "source_simulations", "type": "int", "default": "128", "kind": "keyword_only", "resolved_default": 128 }, { "name": "sequence_length", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "cache_full_shards", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "fields": [], "methods": [ { "name": "__init__", "kind": "method", "signature": "__init__(self, base: ShardedNeuronIODataset, *, window_size: int=500, start_offset: int=500, samples_per_epoch: int | None=None, batch_size: int=8, file_load_fraction: float=0.3, source_simulations: int=128, sequence_length: int | None=None, seed: int=0, cache_full_shards: bool=True) -> None", "parameters": [ { "name": "base", "type": "ShardedNeuronIODataset", "default": null, "kind": "positional_or_keyword" }, { "name": "window_size", "type": "int", "default": "500", "kind": "keyword_only", "resolved_default": 500 }, { "name": "start_offset", "type": "int", "default": "500", "kind": "keyword_only", "resolved_default": 500 }, { "name": "samples_per_epoch", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "batch_size", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "file_load_fraction", "type": "float", "default": "0.3", "kind": "keyword_only", "resolved_default": 0.3 }, { "name": "source_simulations", "type": "int", "default": "128", "kind": "keyword_only", "resolved_default": 128 }, { "name": "sequence_length", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "cache_full_shards", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "line": 405, "end_line": 449, "return_type": "None", "decorators": [] }, { "name": "__len__", "kind": "method", "signature": "__len__(self) -> int", "parameters": [], "line": 451, "end_line": 452, "return_type": "int", "decorators": [] }, { "name": "reshuffle", "kind": "method", "signature": "reshuffle(self, seed: int) -> None", "parameters": [ { "name": "seed", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "line": 454, "end_line": 482, "return_type": "None", "decorators": [] }, { "name": "__getitem__", "kind": "method", "signature": "__getitem__(self, index: int) -> NeuronIOSample", "parameters": [ { "name": "index", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "line": 484, "end_line": 492, "return_type": "NeuronIOSample", "decorators": [] }, { "name": "get_batch", "kind": "method", "signature": "get_batch(self, indices: range | list[int]) -> tuple[np.ndarray, np.ndarray]", "parameters": [ { "name": "indices", "type": "range | list[int]", "default": null, "kind": "positional_or_keyword" } ], "line": 494, "end_line": 511, "return_type": "tuple[np.ndarray, np.ndarray]", "decorators": [] }, { "name": "get_sample_ids", "kind": "method", "signature": "get_sample_ids(self, indices: range | list[int]) -> list[str]", "parameters": [ { "name": "indices", "type": "range | list[int]", "default": null, "kind": "positional_or_keyword" } ], "line": 513, "end_line": 514, "return_type": "list[str]", "decorators": [] } ], "bases": [] }, { "module": "data", "name": "write_demo_shards", "kind": "function", "page": "api/data", "line": 546, "end_line": 571, "signature": "write_demo_shards(root: str | Path, *, shard_count: int=2, samples_per_shard: int=8, time_steps: int=32, input_dim: int=64, seed: int=0) -> None", "parameters": [ { "name": "root", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "shard_count", "type": "int", "default": "2", "kind": "keyword_only", "resolved_default": 2 }, { "name": "samples_per_shard", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "time_steps", "type": "int", "default": "32", "kind": "keyword_only", "resolved_default": 32 }, { "name": "input_dim", "type": "int", "default": "64", "kind": "keyword_only", "resolved_default": 64 }, { "name": "seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 } ], "fields": [], "methods": [] }, { "module": "data", "name": "repack_shards_as_npy", "kind": "function", "page": "api/data", "line": 574, "end_line": 617, "signature": "repack_shards_as_npy(input_root: str | Path, output_root: str | Path, *, input_dtype: np.dtype | type=np.int8) -> list[Path]", "parameters": [ { "name": "input_root", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "output_root", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "input_dtype", "type": "np.dtype | type", "default": "np.int8", "kind": "keyword_only" } ], "fields": [], "methods": [] }, { "module": "neuronio_raw", "name": "RawNeuronIO", "kind": "class", "page": "api/neuronio", "line": 17, "end_line": 21, "signature": "RawNeuronIO(inputs: np.ndarray, spikes: np.ndarray, soma: np.ndarray, metadata: dict[str, object]) -> None", "parameters": [ { "name": "inputs", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" }, { "name": "spikes", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" }, { "name": "soma", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" }, { "name": "metadata", "type": "dict[str, object]", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "inputs", "type": "np.ndarray", "default": null }, { "name": "spikes", "type": "np.ndarray", "default": null }, { "name": "soma", "type": "np.ndarray", "default": null }, { "name": "metadata", "type": "dict[str, object]", "default": null } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "neuronio_raw", "name": "parse_neuronio_pickle", "kind": "function", "page": "api/neuronio", "line": 24, "end_line": 64, "signature": "parse_neuronio_pickle(path: str | Path) -> RawNeuronIO", "parameters": [ { "name": "path", "type": "str | Path", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "neuronio_raw", "name": "convert_neuronio_pickles", "kind": "function", "page": "api/neuronio", "line": 67, "end_line": 132, "signature": "convert_neuronio_pickles(input_path: str | Path, output_dir: str | Path, *, shard_size: int=128, window_size: int | None=None, window_stride: int | None=None, ignore_start: int=0, y_soma_threshold: float=DEFAULT_Y_SOMA_THRESHOLD, y_train_soma_bias: float=DEFAULT_Y_TRAIN_SOMA_BIAS, y_train_soma_scale: float=DEFAULT_Y_TRAIN_SOMA_SCALE) -> list[Path]", "parameters": [ { "name": "input_path", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "output_dir", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "shard_size", "type": "int", "default": "128", "kind": "keyword_only", "resolved_default": 128 }, { "name": "window_size", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "window_stride", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "ignore_start", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "y_soma_threshold", "type": "float", "default": "DEFAULT_Y_SOMA_THRESHOLD", "kind": "keyword_only", "resolved_default": -55.0 }, { "name": "y_train_soma_bias", "type": "float", "default": "DEFAULT_Y_TRAIN_SOMA_BIAS", "kind": "keyword_only", "resolved_default": -67.7 }, { "name": "y_train_soma_scale", "type": "float", "default": "DEFAULT_Y_TRAIN_SOMA_SCALE", "kind": "keyword_only", "resolved_default": 0.1 } ], "fields": [], "methods": [] }, { "module": "neuronio_raw", "name": "create_neuronio_input_type", "kind": "function", "page": "api/neuronio", "line": 135, "end_line": 139, "signature": "create_neuronio_input_type(num_input: int) -> np.ndarray", "parameters": [ { "name": "num_input", "type": "int", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "neuronio_raw", "name": "normalize_soma", "kind": "function", "page": "api/neuronio", "line": 142, "end_line": 150, "signature": "normalize_soma(soma: np.ndarray, threshold: float=DEFAULT_Y_SOMA_THRESHOLD, bias: float=DEFAULT_Y_TRAIN_SOMA_BIAS, scale: float=DEFAULT_Y_TRAIN_SOMA_SCALE) -> np.ndarray", "parameters": [ { "name": "soma", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" }, { "name": "threshold", "type": "float", "default": "DEFAULT_Y_SOMA_THRESHOLD", "kind": "positional_or_keyword", "resolved_default": -55.0 }, { "name": "bias", "type": "float", "default": "DEFAULT_Y_TRAIN_SOMA_BIAS", "kind": "positional_or_keyword", "resolved_default": -67.7 }, { "name": "scale", "type": "float", "default": "DEFAULT_Y_TRAIN_SOMA_SCALE", "kind": "positional_or_keyword", "resolved_default": 0.1 } ], "fields": [], "methods": [] }, { "module": "metrics", "name": "soma_rmse", "kind": "function", "page": "api/metrics", "line": 6, "end_line": 9, "signature": "soma_rmse(prediction: np.ndarray, target: np.ndarray) -> float", "parameters": [ { "name": "prediction", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" }, { "name": "target", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "metrics", "name": "binary_auc", "kind": "function", "page": "api/metrics", "line": 12, "end_line": 34, "signature": "binary_auc(scores: np.ndarray, labels: np.ndarray) -> float", "parameters": [ { "name": "scores", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" }, { "name": "labels", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "metrics", "name": "spike_auc", "kind": "function", "page": "api/metrics", "line": 37, "end_line": 38, "signature": "spike_auc(prediction: np.ndarray, target: np.ndarray) -> float", "parameters": [ { "name": "prediction", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" }, { "name": "target", "type": "np.ndarray", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "evaluate", "name": "evaluate_dataset", "kind": "function", "page": "api/metrics", "line": 17, "end_line": 122, "signature": "evaluate_dataset(model: BranchELM, dataset: ShardedNeuronIODataset, *, batch_size: int=8, device: str='cpu', soma_units: str='millivolts', y_train_soma_scale: float=DEFAULT_Y_TRAIN_SOMA_SCALE, ignore_start: int=0, mask_mode: str='ignore-start', stitch_burn_in: int=150, soma_affine_calibration: bool=False, pin_memory: bool=False, non_blocking: bool=True) -> dict[str, float | int]", "parameters": [ { "name": "model", "type": "BranchELM", "default": null, "kind": "positional_or_keyword" }, { "name": "dataset", "type": "ShardedNeuronIODataset", "default": null, "kind": "positional_or_keyword" }, { "name": "batch_size", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "device", "type": "str", "default": "'cpu'", "kind": "keyword_only", "resolved_default": "cpu" }, { "name": "soma_units", "type": "str", "default": "'millivolts'", "kind": "keyword_only", "resolved_default": "millivolts" }, { "name": "y_train_soma_scale", "type": "float", "default": "DEFAULT_Y_TRAIN_SOMA_SCALE", "kind": "keyword_only", "resolved_default": 0.1 }, { "name": "ignore_start", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "mask_mode", "type": "str", "default": "'ignore-start'", "kind": "keyword_only", "resolved_default": "ignore-start" }, { "name": "stitch_burn_in", "type": "int", "default": "150", "kind": "keyword_only", "resolved_default": 150 }, { "name": "soma_affine_calibration", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "pin_memory", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "non_blocking", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "fields": [], "methods": [] }, { "module": "evaluate", "name": "write_metrics", "kind": "function", "page": "api/metrics", "line": 125, "end_line": 128, "signature": "write_metrics(metrics: dict[str, float | int], output: str | Path) -> None", "parameters": [ { "name": "metrics", "type": "dict[str, float | int]", "default": null, "kind": "positional_or_keyword" }, { "name": "output", "type": "str | Path", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "train", "name": "train_dataset", "kind": "function", "page": "api/training", "line": 21, "end_line": 392, "signature": "train_dataset(model: BranchELM, dataset: ShardedNeuronIODataset, *, epochs: int=1, batch_size: int=8, learning_rate: float=0.0005, burn_in: int=0, device: str='cpu', seed: int=0, validation_dataset: ShardedNeuronIODataset | None=None, validation_batch_size: int | None=None, best_checkpoint_path: str | Path | None=None, validation_soma_units: str='millivolts', validation_metric_ignore_start: int=0, validation_metric_mask_mode: str='ignore-start', validation_metric_stitch_burn_in: int=150, validation_soma_affine_calibration: bool=False, max_train_batches: int | None=None, lr_schedule: str='constant', lr_schedule_steps: int | None=None, optimizer_name: str='adam', weight_decay: float=0.0, l1_lambda: float=0.0, spike_loss_weight: float=0.5, soma_loss_weight: float=0.5, sparse_soma_loss_weight: float=0.0, sparse_soma_high_voltage_quantile: float=0.9, sparse_soma_high_dvdt_quantile: float=0.9, sparse_soma_input_event_quantile: float=0.95, sparse_soma_spike_window: int=5, sparse_soma_post_event_window: int=5, sera_soma_loss_weight: float=0.0, sera_soma_min_weight: float=0.05, sera_soma_relevance_power: float=1.0, soma_slope_loss_weight: float=0.0, grad_clip_norm: float=0.0, update_log_interval: int=0, update_log_path: str | Path | None=None, prefetch_batches: int=0, pin_memory: bool=False, non_blocking: bool=True) -> dict[str, Any]", "parameters": [ { "name": "model", "type": "BranchELM", "default": null, "kind": "positional_or_keyword" }, { "name": "dataset", "type": "ShardedNeuronIODataset", "default": null, "kind": "positional_or_keyword" }, { "name": "epochs", "type": "int", "default": "1", "kind": "keyword_only", "resolved_default": 1 }, { "name": "batch_size", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "learning_rate", "type": "float", "default": "0.0005", "kind": "keyword_only", "resolved_default": 0.0005 }, { "name": "burn_in", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "device", "type": "str", "default": "'cpu'", "kind": "keyword_only", "resolved_default": "cpu" }, { "name": "seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "validation_dataset", "type": "ShardedNeuronIODataset | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "validation_batch_size", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "best_checkpoint_path", "type": "str | Path | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "validation_soma_units", "type": "str", "default": "'millivolts'", "kind": "keyword_only", "resolved_default": "millivolts" }, { "name": "validation_metric_ignore_start", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "validation_metric_mask_mode", "type": "str", "default": "'ignore-start'", "kind": "keyword_only", "resolved_default": "ignore-start" }, { "name": "validation_metric_stitch_burn_in", "type": "int", "default": "150", "kind": "keyword_only", "resolved_default": 150 }, { "name": "validation_soma_affine_calibration", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "max_train_batches", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "lr_schedule", "type": "str", "default": "'constant'", "kind": "keyword_only", "resolved_default": "constant" }, { "name": "lr_schedule_steps", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "optimizer_name", "type": "str", "default": "'adam'", "kind": "keyword_only", "resolved_default": "adam" }, { "name": "weight_decay", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "l1_lambda", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "spike_loss_weight", "type": "float", "default": "0.5", "kind": "keyword_only", "resolved_default": 0.5 }, { "name": "soma_loss_weight", "type": "float", "default": "0.5", "kind": "keyword_only", "resolved_default": 0.5 }, { "name": "sparse_soma_loss_weight", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "sparse_soma_high_voltage_quantile", "type": "float", "default": "0.9", "kind": "keyword_only", "resolved_default": 0.9 }, { "name": "sparse_soma_high_dvdt_quantile", "type": "float", "default": "0.9", "kind": "keyword_only", "resolved_default": 0.9 }, { "name": "sparse_soma_input_event_quantile", "type": "float", "default": "0.95", "kind": "keyword_only", "resolved_default": 0.95 }, { "name": "sparse_soma_spike_window", "type": "int", "default": "5", "kind": "keyword_only", "resolved_default": 5 }, { "name": "sparse_soma_post_event_window", "type": "int", "default": "5", "kind": "keyword_only", "resolved_default": 5 }, { "name": "sera_soma_loss_weight", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "sera_soma_min_weight", "type": "float", "default": "0.05", "kind": "keyword_only", "resolved_default": 0.05 }, { "name": "sera_soma_relevance_power", "type": "float", "default": "1.0", "kind": "keyword_only", "resolved_default": 1.0 }, { "name": "soma_slope_loss_weight", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "grad_clip_norm", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "update_log_interval", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "update_log_path", "type": "str | Path | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "prefetch_batches", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "pin_memory", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "non_blocking", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "fields": [], "methods": [] }, { "module": "train", "name": "train_and_save", "kind": "function", "page": "api/training", "line": 395, "end_line": 493, "signature": "train_and_save(model: BranchELM, dataset: ShardedNeuronIODataset, *, checkpoint_path: str | Path, metrics_path: str | Path, epochs: int=1, batch_size: int=8, learning_rate: float=0.0005, burn_in: int=0, device: str='cpu', seed: int=0, validation_dataset: ShardedNeuronIODataset | None=None, validation_batch_size: int | None=None, best_checkpoint_path: str | Path | None=None, validation_soma_units: str='millivolts', validation_metric_ignore_start: int=0, validation_metric_mask_mode: str='ignore-start', validation_metric_stitch_burn_in: int=150, validation_soma_affine_calibration: bool=False, max_train_batches: int | None=None, lr_schedule: str='constant', lr_schedule_steps: int | None=None, optimizer_name: str='adam', weight_decay: float=0.0, l1_lambda: float=0.0, spike_loss_weight: float=0.5, soma_loss_weight: float=0.5, sparse_soma_loss_weight: float=0.0, sparse_soma_high_voltage_quantile: float=0.9, sparse_soma_high_dvdt_quantile: float=0.9, sparse_soma_input_event_quantile: float=0.95, sparse_soma_spike_window: int=5, sparse_soma_post_event_window: int=5, sera_soma_loss_weight: float=0.0, sera_soma_min_weight: float=0.05, sera_soma_relevance_power: float=1.0, soma_slope_loss_weight: float=0.0, grad_clip_norm: float=0.0, update_log_interval: int=0, update_log_path: str | Path | None=None, prefetch_batches: int=0, pin_memory: bool=False, non_blocking: bool=True) -> dict[str, Any]", "parameters": [ { "name": "model", "type": "BranchELM", "default": null, "kind": "positional_or_keyword" }, { "name": "dataset", "type": "ShardedNeuronIODataset", "default": null, "kind": "positional_or_keyword" }, { "name": "checkpoint_path", "type": "str | Path", "default": null, "kind": "keyword_only" }, { "name": "metrics_path", "type": "str | Path", "default": null, "kind": "keyword_only" }, { "name": "epochs", "type": "int", "default": "1", "kind": "keyword_only", "resolved_default": 1 }, { "name": "batch_size", "type": "int", "default": "8", "kind": "keyword_only", "resolved_default": 8 }, { "name": "learning_rate", "type": "float", "default": "0.0005", "kind": "keyword_only", "resolved_default": 0.0005 }, { "name": "burn_in", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "device", "type": "str", "default": "'cpu'", "kind": "keyword_only", "resolved_default": "cpu" }, { "name": "seed", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "validation_dataset", "type": "ShardedNeuronIODataset | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "validation_batch_size", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "best_checkpoint_path", "type": "str | Path | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "validation_soma_units", "type": "str", "default": "'millivolts'", "kind": "keyword_only", "resolved_default": "millivolts" }, { "name": "validation_metric_ignore_start", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "validation_metric_mask_mode", "type": "str", "default": "'ignore-start'", "kind": "keyword_only", "resolved_default": "ignore-start" }, { "name": "validation_metric_stitch_burn_in", "type": "int", "default": "150", "kind": "keyword_only", "resolved_default": 150 }, { "name": "validation_soma_affine_calibration", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "max_train_batches", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "lr_schedule", "type": "str", "default": "'constant'", "kind": "keyword_only", "resolved_default": "constant" }, { "name": "lr_schedule_steps", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "optimizer_name", "type": "str", "default": "'adam'", "kind": "keyword_only", "resolved_default": "adam" }, { "name": "weight_decay", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "l1_lambda", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "spike_loss_weight", "type": "float", "default": "0.5", "kind": "keyword_only", "resolved_default": 0.5 }, { "name": "soma_loss_weight", "type": "float", "default": "0.5", "kind": "keyword_only", "resolved_default": 0.5 }, { "name": "sparse_soma_loss_weight", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "sparse_soma_high_voltage_quantile", "type": "float", "default": "0.9", "kind": "keyword_only", "resolved_default": 0.9 }, { "name": "sparse_soma_high_dvdt_quantile", "type": "float", "default": "0.9", "kind": "keyword_only", "resolved_default": 0.9 }, { "name": "sparse_soma_input_event_quantile", "type": "float", "default": "0.95", "kind": "keyword_only", "resolved_default": 0.95 }, { "name": "sparse_soma_spike_window", "type": "int", "default": "5", "kind": "keyword_only", "resolved_default": 5 }, { "name": "sparse_soma_post_event_window", "type": "int", "default": "5", "kind": "keyword_only", "resolved_default": 5 }, { "name": "sera_soma_loss_weight", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "sera_soma_min_weight", "type": "float", "default": "0.05", "kind": "keyword_only", "resolved_default": 0.05 }, { "name": "sera_soma_relevance_power", "type": "float", "default": "1.0", "kind": "keyword_only", "resolved_default": 1.0 }, { "name": "soma_slope_loss_weight", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "grad_clip_norm", "type": "float", "default": "0.0", "kind": "keyword_only", "resolved_default": 0.0 }, { "name": "update_log_interval", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "update_log_path", "type": "str | Path | None", "default": "None", "kind": "keyword_only", "resolved_default": null }, { "name": "prefetch_batches", "type": "int", "default": "0", "kind": "keyword_only", "resolved_default": 0 }, { "name": "pin_memory", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "non_blocking", "type": "bool", "default": "True", "kind": "keyword_only", "resolved_default": true } ], "fields": [], "methods": [] }, { "module": "inference_benchmark", "name": "benchmark_inference_matrix", "kind": "function", "page": "api/inference-benchmark", "line": 12, "end_line": 103, "signature": "benchmark_inference_matrix(model: torch.nn.Module, *, batch_sizes: Iterable[int], time_steps: Iterable[int], input_dim: int, device: str='cpu', precision: str='float32', warmup_runs: int=5, runs: int=20, compile_model: bool=False, accuracy_metrics_path: str | Path | None=None) -> dict[str, Any]", "parameters": [ { "name": "model", "type": "torch.nn.Module", "default": null, "kind": "positional_or_keyword" }, { "name": "batch_sizes", "type": "Iterable[int]", "default": null, "kind": "keyword_only" }, { "name": "time_steps", "type": "Iterable[int]", "default": null, "kind": "keyword_only" }, { "name": "input_dim", "type": "int", "default": null, "kind": "keyword_only" }, { "name": "device", "type": "str", "default": "'cpu'", "kind": "keyword_only", "resolved_default": "cpu" }, { "name": "precision", "type": "str", "default": "'float32'", "kind": "keyword_only", "resolved_default": "float32" }, { "name": "warmup_runs", "type": "int", "default": "5", "kind": "keyword_only", "resolved_default": 5 }, { "name": "runs", "type": "int", "default": "20", "kind": "keyword_only", "resolved_default": 20 }, { "name": "compile_model", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "accuracy_metrics_path", "type": "str | Path | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "fields": [], "methods": [] }, { "module": "inference_benchmark", "name": "count_parameters", "kind": "function", "page": "api/inference-benchmark", "line": 106, "end_line": 107, "signature": "count_parameters(model: torch.nn.Module) -> int", "parameters": [ { "name": "model", "type": "torch.nn.Module", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "inference_benchmark", "name": "write_inference_benchmark", "kind": "function", "page": "api/inference-benchmark", "line": 110, "end_line": 113, "signature": "write_inference_benchmark(report: dict[str, Any], path: str | Path) -> None", "parameters": [ { "name": "report", "type": "dict[str, Any]", "default": null, "kind": "positional_or_keyword" }, { "name": "path", "type": "str | Path", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "inference_benchmark", "name": "parse_int_list", "kind": "function", "page": "api/inference-benchmark", "line": 116, "end_line": 120, "signature": "parse_int_list(value: str) -> list[int]", "parameters": [ { "name": "value", "type": "str", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "axobench_iteration", "name": "make_official_elm_predictor", "kind": "function", "page": "api/axobench", "line": 22, "end_line": 83, "signature": "make_official_elm_predictor(config_path: str | Path, checkpoint_path: str | Path, *, device: str='auto', dtype: str | torch.dtype='float32') -> tuple[Callable[[np.ndarray], np.ndarray], dict[str, Any]]", "parameters": [ { "name": "config_path", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "checkpoint_path", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "device", "type": "str", "default": "'auto'", "kind": "keyword_only", "resolved_default": "auto" }, { "name": "dtype", "type": "str | torch.dtype", "default": "'float32'", "kind": "keyword_only", "resolved_default": "float32" } ], "fields": [], "methods": [] }, { "module": "axobench_iteration", "name": "make_checkpoint_predictor", "kind": "function", "page": "api/axobench", "line": 86, "end_line": 189, "signature": "make_checkpoint_predictor(checkpoint: str | Path, *, device: str='auto', dtype: str | torch.dtype='float32', streaming: bool=False, streaming_chunk_size: int | None=None) -> tuple[Callable[[np.ndarray], np.ndarray], dict[str, Any]]", "parameters": [ { "name": "checkpoint", "type": "str | Path", "default": null, "kind": "positional_or_keyword" }, { "name": "device", "type": "str", "default": "'auto'", "kind": "keyword_only", "resolved_default": "auto" }, { "name": "dtype", "type": "str | torch.dtype", "default": "'float32'", "kind": "keyword_only", "resolved_default": "float32" }, { "name": "streaming", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false }, { "name": "streaming_chunk_size", "type": "int | None", "default": "None", "kind": "keyword_only", "resolved_default": null } ], "fields": [], "methods": [] }, { "module": "setup_workflow", "name": "DatasetSpec", "kind": "class", "page": "api/setup", "line": 12, "end_line": 14, "signature": "DatasetSpec(name: str, slug: str) -> None", "parameters": [ { "name": "name", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "slug", "type": "str", "default": null, "kind": "positional_or_keyword" } ], "fields": [ { "name": "name", "type": "str", "default": null }, { "name": "slug", "type": "str", "default": null } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "setup_workflow", "name": "SetupOptions", "kind": "class", "page": "api/setup", "line": 24, "end_line": 35, "signature": "SetupOptions(project_root: Path, data_dir: Path, install_kaggle: bool = False, download_data: bool = False, convert_raw: bool = False, raw_dir: Path | None = None, shard_dir: Path | None = None, shard_size: int = 128, skip_install: bool = False, datasets: tuple[DatasetSpec, ...] = DEFAULT_DATASETS, python_executable: str = sys.executable) -> None", "parameters": [ { "name": "project_root", "type": "Path", "default": null, "kind": "positional_or_keyword" }, { "name": "data_dir", "type": "Path", "default": null, "kind": "positional_or_keyword" }, { "name": "install_kaggle", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "download_data", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "convert_raw", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "raw_dir", "type": "Path | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "shard_dir", "type": "Path | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "shard_size", "type": "int", "default": "128", "kind": "positional_or_keyword" }, { "name": "skip_install", "type": "bool", "default": "False", "kind": "positional_or_keyword" }, { "name": "datasets", "type": "tuple[DatasetSpec, ...]", "default": "DEFAULT_DATASETS", "kind": "positional_or_keyword" }, { "name": "python_executable", "type": "str", "default": "sys.executable", "kind": "positional_or_keyword" } ], "fields": [ { "name": "project_root", "type": "Path", "default": null }, { "name": "data_dir", "type": "Path", "default": null }, { "name": "install_kaggle", "type": "bool", "default": "False" }, { "name": "download_data", "type": "bool", "default": "False" }, { "name": "convert_raw", "type": "bool", "default": "False" }, { "name": "raw_dir", "type": "Path | None", "default": "None" }, { "name": "shard_dir", "type": "Path | None", "default": "None" }, { "name": "shard_size", "type": "int", "default": "128" }, { "name": "skip_install", "type": "bool", "default": "False" }, { "name": "datasets", "type": "tuple[DatasetSpec, ...]", "default": "DEFAULT_DATASETS" }, { "name": "python_executable", "type": "str", "default": "sys.executable" } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "setup_workflow", "name": "SetupStep", "kind": "class", "page": "api/setup", "line": 39, "end_line": 43, "signature": "SetupStep(description: str, command: str | None = None, argv: tuple[str, ...] | None = None, mkdir: Path | None = None) -> None", "parameters": [ { "name": "description", "type": "str", "default": null, "kind": "positional_or_keyword" }, { "name": "command", "type": "str | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "argv", "type": "tuple[str, ...] | None", "default": "None", "kind": "positional_or_keyword" }, { "name": "mkdir", "type": "Path | None", "default": "None", "kind": "positional_or_keyword" } ], "fields": [ { "name": "description", "type": "str", "default": null }, { "name": "command", "type": "str | None", "default": "None" }, { "name": "argv", "type": "tuple[str, ...] | None", "default": "None" }, { "name": "mkdir", "type": "Path | None", "default": "None" } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "setup_workflow", "name": "SetupPlan", "kind": "class", "page": "api/setup", "line": 47, "end_line": 49, "signature": "SetupPlan(steps: tuple[SetupStep, ...], next_steps: tuple[str, ...] = field(default_factory=tuple)) -> None", "parameters": [ { "name": "steps", "type": "tuple[SetupStep, ...]", "default": null, "kind": "positional_or_keyword" }, { "name": "next_steps", "type": "tuple[str, ...]", "default": "field(default_factory=tuple)", "kind": "positional_or_keyword" } ], "fields": [ { "name": "steps", "type": "tuple[SetupStep, ...]", "default": null }, { "name": "next_steps", "type": "tuple[str, ...]", "default": "field(default_factory=tuple)" } ], "methods": [], "bases": [], "inherited_fields": [] }, { "module": "setup_workflow", "name": "build_setup_plan", "kind": "function", "page": "api/setup", "line": 52, "end_line": 128, "signature": "build_setup_plan(options: SetupOptions) -> SetupPlan", "parameters": [ { "name": "options", "type": "SetupOptions", "default": null, "kind": "positional_or_keyword" } ], "fields": [], "methods": [] }, { "module": "setup_workflow", "name": "run_setup_plan", "kind": "function", "page": "api/setup", "line": 131, "end_line": 145, "signature": "run_setup_plan(plan: SetupPlan, *, dry_run: bool=False) -> None", "parameters": [ { "name": "plan", "type": "SetupPlan", "default": null, "kind": "positional_or_keyword" }, { "name": "dry_run", "type": "bool", "default": "False", "kind": "keyword_only", "resolved_default": false } ], "fields": [], "methods": [] } ], "cli": [ { "command": "axosim", "source_module": "cli", "sections": [ { "name": "axosim", "options": [], "description": "" }, { "name": "axosim make-demo-data", "options": [ { "flags": [ "--output" ], "dest": "output", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 61, "hidden": false }, { "flags": [ "--samples" ], "dest": "samples", "type": "int", "default": 8, "required": false, "choices": null, "action": "store", "help": "", "line": 62, "hidden": false }, { "flags": [ "--shards" ], "dest": "shards", "type": "int", "default": 1, "required": false, "choices": null, "action": "store", "help": "", "line": 63, "hidden": false }, { "flags": [ "--time-steps" ], "dest": "time_steps", "type": "int", "default": 32, "required": false, "choices": null, "action": "store", "help": "", "line": 64, "hidden": false }, { "flags": [ "--input-dim" ], "dest": "input_dim", "type": "int", "default": 64, "required": false, "choices": null, "action": "store", "help": "", "line": 65, "hidden": false }, { "flags": [ "--seed" ], "dest": "seed", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 66, "hidden": false } ], "description": "write deterministic demo NPZ shards" }, { "name": "axosim convert-neuronio", "options": [ { "flags": [ "--input" ], "dest": "input", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 69, "hidden": false }, { "flags": [ "--output" ], "dest": "output", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 70, "hidden": false }, { "flags": [ "--shard-size" ], "dest": "shard_size", "type": "int", "default": 128, "required": false, "choices": null, "action": "store", "help": "", "line": 71, "hidden": false }, { "flags": [ "--window-size" ], "dest": "window_size", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 72, "hidden": false }, { "flags": [ "--window-stride" ], "dest": "window_stride", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 73, "hidden": false }, { "flags": [ "--ignore-start" ], "dest": "ignore_start", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 74, "hidden": false } ], "description": "convert raw NeuronIO pickle files to deterministic shards" }, { "name": "axosim repack-shards", "options": [ { "flags": [ "--input" ], "dest": "input", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 77, "hidden": false }, { "flags": [ "--output" ], "dest": "output", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 78, "hidden": false }, { "flags": [ "--input-dtype" ], "dest": "input_dtype", "type": "str", "default": "int8", "required": false, "choices": [ "int8", "float32" ], "action": "store", "help": "", "line": 79, "hidden": false } ], "description": "repack shards as sliceable .npy arrays" }, { "name": "axosim list-presets", "options": [], "description": "list named train/evaluate presets" }, { "name": "axosim evaluate", "options": [ { "flags": [ "--preset" ], "dest": "preset", "type": "str", "default": null, "required": false, "choices": [ "fulltrace-paper-metric", "legacy-baseline-smoke" ], "action": "store", "help": "", "line": 84, "hidden": false }, { "flags": [ "--data" ], "dest": "data", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 85, "hidden": false }, { "flags": [ "--output" ], "dest": "output", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 86, "hidden": false }, { "flags": [ "--input-dim" ], "dest": "input_dim", "type": "int", "default": 1278, "required": false, "choices": null, "action": "store", "help": "", "line": 87, "hidden": false }, { "flags": [ "--memory-units" ], "dest": "memory_units", "type": "int", "default": 30, "required": false, "choices": null, "action": "store", "help": "", "line": 88, "hidden": false }, { "flags": [ "--branches" ], "dest": "branches", "type": "int", "default": 32, "required": false, "choices": null, "action": "store", "help": "", "line": 89, "hidden": false }, { "flags": [ "--model-kind" ], "dest": "model_kind", "type": "str", "default": "baseline", "required": false, "choices": [ "baseline", "official", "axomamba", "mamba-official", "mamba-pytorch", "branch-mamba-pytorch", "branch-trace-rnn" ], "action": "store", "help": "use 'official' for the paper-style model, 'axomamba' for the promoted Mamba surrogate, 'mamba-official' for the generic official mamba-ssm backend, or 'branch-trace-rnn' for the RNN comparison path; 'baseline' is legacy/debug", "line": 90, "hidden": false }, { "flags": [ "--model-config" ], "dest": "model_config", "type": "str", "default": "configs/branch_elm_30_official.json", "required": false, "choices": null, "action": "store", "help": "", "line": 91, "hidden": false }, { "flags": [ "--batch-size" ], "dest": "batch_size", "type": "int", "default": 8, "required": false, "choices": null, "action": "store", "help": "", "line": 92, "hidden": false }, { "flags": [ "--seed" ], "dest": "seed", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 93, "hidden": false }, { "flags": [ "--device" ], "dest": "device", "type": "str", "default": "cpu", "required": false, "choices": null, "action": "store", "help": "", "line": 94, "hidden": false }, { "flags": [ "--checkpoint" ], "dest": "checkpoint", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 95, "hidden": false }, { "flags": [ "--window-size" ], "dest": "window_size", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 96, "hidden": false }, { "flags": [ "--window-stride" ], "dest": "window_stride", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 97, "hidden": false }, { "flags": [ "--ignore-start" ], "dest": "ignore_start", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 98, "hidden": false }, { "flags": [ "--cache-shards" ], "dest": "cache_shards", "type": "int", "default": 1, "required": false, "choices": null, "action": "store", "help": "", "line": 99, "hidden": false }, { "flags": [ "--soma-units" ], "dest": "soma_units", "type": "str", "default": "millivolts", "required": false, "choices": [ "millivolts", "normalized" ], "action": "store", "help": "", "line": 100, "hidden": false }, { "flags": [ "--metric-ignore-start" ], "dest": "metric_ignore_start", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 101, "hidden": false }, { "flags": [ "--metric-mask-mode" ], "dest": "metric_mask_mode", "type": "str", "default": "ignore-start", "required": false, "choices": [ "ignore-start", "official-overlap" ], "action": "store", "help": "argparse.SUPPRESS", "line": 102, "hidden": true }, { "flags": [ "--metric-stitch-burn-in" ], "dest": "metric_stitch_burn_in", "type": "int", "default": 150, "required": false, "choices": null, "action": "store", "help": "argparse.SUPPRESS", "line": 103, "hidden": true }, { "flags": [ "--soma-affine-calibration" ], "dest": "soma_affine_calibration", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 104, "hidden": false }, { "flags": [ "--pin-memory" ], "dest": "pin_memory", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 105, "hidden": false }, { "flags": [ "--blocking-transfer" ], "dest": "non_blocking", "type": "flag", "default": true, "required": false, "choices": null, "action": "store_false", "help": "", "line": 106, "hidden": false }, { "flags": [ "--registry" ], "dest": "registry", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "append a one-line JSONL record; defaults to /experiment-registry.jsonl", "line": 107, "hidden": false }, { "flags": [ "--no-registry" ], "dest": "registry", "type": "flag", "default": null, "required": false, "choices": null, "action": "store_const", "help": "", "line": 108, "hidden": false, "const": "off" } ], "description": "evaluate a model with the corrected full-trace metric path" }, { "name": "axosim diagnose", "options": [ { "flags": [ "--preset" ], "dest": "preset", "type": "str", "default": null, "required": false, "choices": [ "fulltrace-paper-metric", "legacy-baseline-smoke" ], "action": "store", "help": "", "line": 112, "hidden": false }, { "flags": [ "--data" ], "dest": "data", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 113, "hidden": false }, { "flags": [ "--output-dir" ], "dest": "output_dir", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 114, "hidden": false }, { "flags": [ "--input-dim" ], "dest": "input_dim", "type": "int", "default": 1278, "required": false, "choices": null, "action": "store", "help": "", "line": 115, "hidden": false }, { "flags": [ "--memory-units" ], "dest": "memory_units", "type": "int", "default": 30, "required": false, "choices": null, "action": "store", "help": "", "line": 116, "hidden": false }, { "flags": [ "--branches" ], "dest": "branches", "type": "int", "default": 32, "required": false, "choices": null, "action": "store", "help": "", "line": 117, "hidden": false }, { "flags": [ "--model-kind" ], "dest": "model_kind", "type": "str", "default": "baseline", "required": false, "choices": [ "baseline", "official", "axomamba", "mamba-official", "mamba-pytorch", "branch-mamba-pytorch", "branch-trace-rnn" ], "action": "store", "help": "", "line": 118, "hidden": false }, { "flags": [ "--model-config" ], "dest": "model_config", "type": "str", "default": "configs/branch_elm_30_official.json", "required": false, "choices": null, "action": "store", "help": "", "line": 119, "hidden": false }, { "flags": [ "--batch-size" ], "dest": "batch_size", "type": "int", "default": 8, "required": false, "choices": null, "action": "store", "help": "", "line": 120, "hidden": false }, { "flags": [ "--seed" ], "dest": "seed", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 121, "hidden": false }, { "flags": [ "--device" ], "dest": "device", "type": "str", "default": "cpu", "required": false, "choices": null, "action": "store", "help": "", "line": 122, "hidden": false }, { "flags": [ "--checkpoint" ], "dest": "checkpoint", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 123, "hidden": false }, { "flags": [ "--window-size" ], "dest": "window_size", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 124, "hidden": false }, { "flags": [ "--window-stride" ], "dest": "window_stride", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 125, "hidden": false }, { "flags": [ "--ignore-start" ], "dest": "ignore_start", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 126, "hidden": false }, { "flags": [ "--cache-shards" ], "dest": "cache_shards", "type": "int", "default": 1, "required": false, "choices": null, "action": "store", "help": "", "line": 127, "hidden": false }, { "flags": [ "--metric-ignore-start" ], "dest": "metric_ignore_start", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 128, "hidden": false }, { "flags": [ "--metric-mask-mode" ], "dest": "metric_mask_mode", "type": "str", "default": "ignore-start", "required": false, "choices": [ "ignore-start", "official-overlap" ], "action": "store", "help": "argparse.SUPPRESS", "line": 129, "hidden": true }, { "flags": [ "--metric-stitch-burn-in" ], "dest": "metric_stitch_burn_in", "type": "int", "default": 150, "required": false, "choices": null, "action": "store", "help": "argparse.SUPPRESS", "line": 130, "hidden": true }, { "flags": [ "--soma-affine-calibration" ], "dest": "soma_affine_calibration", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 131, "hidden": false }, { "flags": [ "--max-samples" ], "dest": "max_samples", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 132, "hidden": false }, { "flags": [ "--pin-memory" ], "dest": "pin_memory", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 133, "hidden": false }, { "flags": [ "--blocking-transfer" ], "dest": "non_blocking", "type": "flag", "default": true, "required": false, "choices": null, "action": "store_false", "help": "", "line": 134, "hidden": false } ], "description": "write biology-oriented surrogate fidelity diagnostics" }, { "name": "axosim train", "options": [ { "flags": [ "--preset" ], "dest": "preset", "type": "str", "default": null, "required": false, "choices": [ "axomamba-probe", "legacy-baseline-smoke", "mamba-official-probe", "mamba2-official-probe", "official-fulltrace-reference", "paper-random-fulltrace" ], "action": "store", "help": "", "line": 138, "hidden": false }, { "flags": [ "--data" ], "dest": "data", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 139, "hidden": false }, { "flags": [ "--checkpoint" ], "dest": "checkpoint", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 140, "hidden": false }, { "flags": [ "--metrics" ], "dest": "metrics", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 141, "hidden": false }, { "flags": [ "--input-dim" ], "dest": "input_dim", "type": "int", "default": 1278, "required": false, "choices": null, "action": "store", "help": "", "line": 142, "hidden": false }, { "flags": [ "--memory-units" ], "dest": "memory_units", "type": "int", "default": 30, "required": false, "choices": null, "action": "store", "help": "", "line": 143, "hidden": false }, { "flags": [ "--branches" ], "dest": "branches", "type": "int", "default": 32, "required": false, "choices": null, "action": "store", "help": "", "line": 144, "hidden": false }, { "flags": [ "--model-kind" ], "dest": "model_kind", "type": "str", "default": "baseline", "required": false, "choices": [ "baseline", "official", "axomamba", "mamba-official", "mamba-pytorch", "branch-mamba-pytorch", "branch-trace-rnn" ], "action": "store", "help": "use 'official' for the paper-style model, 'axomamba' for the promoted Mamba surrogate, 'mamba-official' for the generic official mamba-ssm backend, or 'branch-trace-rnn' for the RNN comparison path; 'baseline' is legacy/debug", "line": 145, "hidden": false }, { "flags": [ "--model-config" ], "dest": "model_config", "type": "str", "default": "configs/branch_elm_30_official.json", "required": false, "choices": null, "action": "store", "help": "", "line": 146, "hidden": false }, { "flags": [ "--epochs" ], "dest": "epochs", "type": "int", "default": 1, "required": false, "choices": null, "action": "store", "help": "", "line": 147, "hidden": false }, { "flags": [ "--batch-size" ], "dest": "batch_size", "type": "int", "default": 8, "required": false, "choices": null, "action": "store", "help": "", "line": 148, "hidden": false }, { "flags": [ "--learning-rate" ], "dest": "learning_rate", "type": "float", "default": 0.0005, "required": false, "choices": null, "action": "store", "help": "", "line": 149, "hidden": false }, { "flags": [ "--burn-in" ], "dest": "burn_in", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 150, "hidden": false }, { "flags": [ "--seed" ], "dest": "seed", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 151, "hidden": false }, { "flags": [ "--device" ], "dest": "device", "type": "str", "default": "cpu", "required": false, "choices": null, "action": "store", "help": "", "line": 152, "hidden": false }, { "flags": [ "--window-size" ], "dest": "window_size", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 153, "hidden": false }, { "flags": [ "--window-stride" ], "dest": "window_stride", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 154, "hidden": false }, { "flags": [ "--ignore-start" ], "dest": "ignore_start", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 155, "hidden": false }, { "flags": [ "--window-sampling" ], "dest": "window_sampling", "type": "str", "default": "random-full-trace", "required": false, "choices": [ "deterministic", "random-full-trace", "official-full-trace" ], "action": "store", "help": "", "line": 156, "hidden": false }, { "flags": [ "--epoch-samples" ], "dest": "epoch_samples", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 161, "hidden": false }, { "flags": [ "--shard-reuse-batches" ], "dest": "shard_reuse_batches", "type": "int", "default": 1, "required": false, "choices": null, "action": "store", "help": "", "line": 162, "hidden": false }, { "flags": [ "--file-load-fraction" ], "dest": "file_load_fraction", "type": "float", "default": 0.3, "required": false, "choices": null, "action": "store", "help": "", "line": 163, "hidden": false }, { "flags": [ "--full-trace-length" ], "dest": "full_trace_length", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 164, "hidden": false }, { "flags": [ "--sample-window-reads" ], "dest": "sample_window_reads", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "read random/official windows by sample instead of caching full compressed shards", "line": 165, "hidden": false }, { "flags": [ "--shuffle-mode" ], "dest": "shuffle_mode", "type": "str", "default": "shard", "required": false, "choices": [ "sample", "shard", "none" ], "action": "store", "help": "", "line": 170, "hidden": false }, { "flags": [ "--cache-shards" ], "dest": "cache_shards", "type": "int", "default": 1, "required": false, "choices": null, "action": "store", "help": "", "line": 171, "hidden": false }, { "flags": [ "--val-data" ], "dest": "val_data", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 172, "hidden": false }, { "flags": [ "--val-batch-size" ], "dest": "val_batch_size", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 173, "hidden": false }, { "flags": [ "--val-cache-shards" ], "dest": "val_cache_shards", "type": "int", "default": 1, "required": false, "choices": null, "action": "store", "help": "", "line": 174, "hidden": false }, { "flags": [ "--val-window-size" ], "dest": "val_window_size", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 175, "hidden": false }, { "flags": [ "--val-window-stride" ], "dest": "val_window_stride", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 176, "hidden": false }, { "flags": [ "--val-ignore-start" ], "dest": "val_ignore_start", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 177, "hidden": false }, { "flags": [ "--val-soma-units" ], "dest": "val_soma_units", "type": "str", "default": "millivolts", "required": false, "choices": [ "millivolts", "normalized" ], "action": "store", "help": "", "line": 178, "hidden": false }, { "flags": [ "--val-metric-ignore-start" ], "dest": "val_metric_ignore_start", "type": "int", "default": 500, "required": false, "choices": null, "action": "store", "help": "", "line": 179, "hidden": false }, { "flags": [ "--val-metric-mask-mode" ], "dest": "val_metric_mask_mode", "type": "str", "default": "ignore-start", "required": false, "choices": [ "ignore-start", "official-overlap" ], "action": "store", "help": "argparse.SUPPRESS", "line": 180, "hidden": true }, { "flags": [ "--val-metric-stitch-burn-in" ], "dest": "val_metric_stitch_burn_in", "type": "int", "default": 150, "required": false, "choices": null, "action": "store", "help": "argparse.SUPPRESS", "line": 181, "hidden": true }, { "flags": [ "--val-soma-affine-calibration" ], "dest": "val_soma_affine_calibration", "type": "flag", "default": true, "required": false, "choices": null, "action": "store_true", "help": "", "line": 182, "hidden": false }, { "flags": [ "--no-val-soma-affine-calibration" ], "dest": "val_soma_affine_calibration", "type": "flag", "default": true, "required": false, "choices": null, "action": "store_false", "help": "", "line": 183, "hidden": false }, { "flags": [ "--best-checkpoint" ], "dest": "best_checkpoint", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 184, "hidden": false }, { "flags": [ "--max-train-batches" ], "dest": "max_train_batches", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 185, "hidden": false }, { "flags": [ "--lr-schedule" ], "dest": "lr_schedule", "type": "str", "default": "constant", "required": false, "choices": [ "constant", "cosine" ], "action": "store", "help": "", "line": 186, "hidden": false }, { "flags": [ "--lr-schedule-steps" ], "dest": "lr_schedule_steps", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 187, "hidden": false }, { "flags": [ "--optimizer" ], "dest": "optimizer", "type": "str", "default": "adam", "required": false, "choices": [ "adam", "adamw" ], "action": "store", "help": "", "line": 188, "hidden": false }, { "flags": [ "--weight-decay" ], "dest": "weight_decay", "type": "float", "default": 0.0, "required": false, "choices": null, "action": "store", "help": "", "line": 189, "hidden": false }, { "flags": [ "--l1-lambda" ], "dest": "l1_lambda", "type": "float", "default": 0.0, "required": false, "choices": null, "action": "store", "help": "", "line": 190, "hidden": false }, { "flags": [ "--spike-loss-weight" ], "dest": "spike_loss_weight", "type": "float", "default": 0.5, "required": false, "choices": null, "action": "store", "help": "", "line": 191, "hidden": false }, { "flags": [ "--soma-loss-weight" ], "dest": "soma_loss_weight", "type": "float", "default": 0.5, "required": false, "choices": null, "action": "store", "help": "", "line": 192, "hidden": false }, { "flags": [ "--sparse-soma-loss-weight" ], "dest": "sparse_soma_loss_weight", "type": "float", "default": 0.0, "required": false, "choices": null, "action": "store", "help": "", "line": 193, "hidden": false }, { "flags": [ "--sparse-soma-high-voltage-quantile" ], "dest": "sparse_soma_high_voltage_quantile", "type": "float", "default": 0.9, "required": false, "choices": null, "action": "store", "help": "", "line": 194, "hidden": false }, { "flags": [ "--sparse-soma-high-dvdt-quantile" ], "dest": "sparse_soma_high_dvdt_quantile", "type": "float", "default": 0.9, "required": false, "choices": null, "action": "store", "help": "", "line": 195, "hidden": false }, { "flags": [ "--sparse-soma-input-event-quantile" ], "dest": "sparse_soma_input_event_quantile", "type": "float", "default": 0.95, "required": false, "choices": null, "action": "store", "help": "", "line": 196, "hidden": false }, { "flags": [ "--sparse-soma-spike-window" ], "dest": "sparse_soma_spike_window", "type": "int", "default": 5, "required": false, "choices": null, "action": "store", "help": "", "line": 197, "hidden": false }, { "flags": [ "--sparse-soma-post-event-window" ], "dest": "sparse_soma_post_event_window", "type": "int", "default": 5, "required": false, "choices": null, "action": "store", "help": "", "line": 198, "hidden": false }, { "flags": [ "--sera-soma-loss-weight" ], "dest": "sera_soma_loss_weight", "type": "float", "default": 0.0, "required": false, "choices": null, "action": "store", "help": "", "line": 199, "hidden": false }, { "flags": [ "--sera-soma-min-weight" ], "dest": "sera_soma_min_weight", "type": "float", "default": 0.05, "required": false, "choices": null, "action": "store", "help": "", "line": 200, "hidden": false }, { "flags": [ "--sera-soma-relevance-power" ], "dest": "sera_soma_relevance_power", "type": "float", "default": 1.0, "required": false, "choices": null, "action": "store", "help": "", "line": 201, "hidden": false }, { "flags": [ "--soma-slope-loss-weight" ], "dest": "soma_slope_loss_weight", "type": "float", "default": 0.0, "required": false, "choices": null, "action": "store", "help": "", "line": 202, "hidden": false }, { "flags": [ "--grad-clip-norm" ], "dest": "grad_clip_norm", "type": "float", "default": 0.0, "required": false, "choices": null, "action": "store", "help": "", "line": 203, "hidden": false }, { "flags": [ "--train-update-log-interval" ], "dest": "train_update_log_interval", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 204, "hidden": false }, { "flags": [ "--train-update-log" ], "dest": "train_update_log", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 205, "hidden": false }, { "flags": [ "--prefetch-batches" ], "dest": "prefetch_batches", "type": "int", "default": 0, "required": false, "choices": null, "action": "store", "help": "", "line": 206, "hidden": false }, { "flags": [ "--pin-memory" ], "dest": "pin_memory", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 207, "hidden": false }, { "flags": [ "--blocking-transfer" ], "dest": "non_blocking", "type": "flag", "default": true, "required": false, "choices": null, "action": "store_false", "help": "", "line": 208, "hidden": false }, { "flags": [ "--registry" ], "dest": "registry", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "append a one-line JSONL record; defaults to /experiment-registry.jsonl", "line": 209, "hidden": false }, { "flags": [ "--no-registry" ], "dest": "registry", "type": "flag", "default": null, "required": false, "choices": null, "action": "store_const", "help": "", "line": 210, "hidden": false, "const": "off" } ], "description": "train a model with the current experiment workflow" }, { "name": "axosim benchmark", "options": [ { "flags": [ "--input-dim" ], "dest": "input_dim", "type": "int", "default": 1278, "required": false, "choices": null, "action": "store", "help": "", "line": 214, "hidden": false }, { "flags": [ "--memory-units" ], "dest": "memory_units", "type": "int", "default": 30, "required": false, "choices": null, "action": "store", "help": "", "line": 215, "hidden": false }, { "flags": [ "--branches" ], "dest": "branches", "type": "int", "default": 32, "required": false, "choices": null, "action": "store", "help": "", "line": 216, "hidden": false }, { "flags": [ "--model-kind" ], "dest": "model_kind", "type": "str", "default": "baseline", "required": false, "choices": [ "baseline", "official", "axomamba", "mamba-official", "mamba-pytorch", "branch-mamba-pytorch", "branch-trace-rnn" ], "action": "store", "help": "use 'official' for the paper-style model, 'mamba-official' for the official mamba-ssm backend, or 'branch-trace-rnn' for the RNN comparison path; 'baseline' is legacy/debug", "line": 217, "hidden": false }, { "flags": [ "--model-config" ], "dest": "model_config", "type": "str", "default": "configs/branch_elm_30_official.json", "required": false, "choices": null, "action": "store", "help": "", "line": 218, "hidden": false }, { "flags": [ "--batch-size" ], "dest": "batch_size", "type": "int", "default": 1, "required": false, "choices": null, "action": "store", "help": "", "line": 219, "hidden": false }, { "flags": [ "--time-steps" ], "dest": "time_steps", "type": "int", "default": 500, "required": false, "choices": null, "action": "store", "help": "", "line": 220, "hidden": false }, { "flags": [ "--runs" ], "dest": "runs", "type": "int", "default": 10, "required": false, "choices": null, "action": "store", "help": "", "line": 221, "hidden": false }, { "flags": [ "--device" ], "dest": "device", "type": "str", "default": "cpu", "required": false, "choices": null, "action": "store", "help": "", "line": 222, "hidden": false }, { "flags": [ "--compile" ], "dest": "compile_model", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 223, "hidden": false }, { "flags": [ "--output" ], "dest": "output", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 224, "hidden": false } ], "description": "time model forward passes on random input" }, { "name": "axosim inference-benchmark", "options": [ { "flags": [ "--input-dim" ], "dest": "input_dim", "type": "int", "default": 1278, "required": false, "choices": null, "action": "store", "help": "", "line": 230, "hidden": false }, { "flags": [ "--memory-units" ], "dest": "memory_units", "type": "int", "default": 30, "required": false, "choices": null, "action": "store", "help": "", "line": 231, "hidden": false }, { "flags": [ "--branches" ], "dest": "branches", "type": "int", "default": 32, "required": false, "choices": null, "action": "store", "help": "", "line": 232, "hidden": false }, { "flags": [ "--model-kind" ], "dest": "model_kind", "type": "str", "default": "baseline", "required": false, "choices": [ "baseline", "official", "axomamba", "mamba-official", "mamba-pytorch", "branch-mamba-pytorch", "branch-trace-rnn" ], "action": "store", "help": "", "line": 233, "hidden": false }, { "flags": [ "--model-config" ], "dest": "model_config", "type": "str", "default": "configs/branch_elm_30_official.json", "required": false, "choices": null, "action": "store", "help": "", "line": 234, "hidden": false }, { "flags": [ "--checkpoint" ], "dest": "checkpoint", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 235, "hidden": false }, { "flags": [ "--batch-sizes" ], "dest": "batch_sizes", "type": "str", "default": "1,8,32,128", "required": false, "choices": null, "action": "store", "help": "", "line": 236, "hidden": false }, { "flags": [ "--time-steps" ], "dest": "time_steps", "type": "str", "default": "500,1000,6000", "required": false, "choices": null, "action": "store", "help": "", "line": 237, "hidden": false }, { "flags": [ "--runs" ], "dest": "runs", "type": "int", "default": 20, "required": false, "choices": null, "action": "store", "help": "", "line": 238, "hidden": false }, { "flags": [ "--warmup-runs" ], "dest": "warmup_runs", "type": "int", "default": 5, "required": false, "choices": null, "action": "store", "help": "", "line": 239, "hidden": false }, { "flags": [ "--precision" ], "dest": "precision", "type": "str", "default": "float32", "required": false, "choices": [ "float32", "float16", "bfloat16" ], "action": "store", "help": "", "line": 240, "hidden": false }, { "flags": [ "--device" ], "dest": "device", "type": "str", "default": "cpu", "required": false, "choices": null, "action": "store", "help": "", "line": 241, "hidden": false }, { "flags": [ "--compile" ], "dest": "compile_model", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 242, "hidden": false }, { "flags": [ "--accuracy-metrics" ], "dest": "accuracy_metrics", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 243, "hidden": false }, { "flags": [ "--output" ], "dest": "output", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 244, "hidden": false } ], "description": "benchmark deployment-oriented inference throughput across batch and sequence scales" } ] }, { "command": "axosim-setup", "source_module": "setup_workflow", "sections": [ { "name": "axosim-setup", "options": [ { "flags": [ "--project-root" ], "dest": "project_root", "type": "Path", "default": "Path.cwd()", "required": false, "choices": null, "action": "store", "help": "", "line": 152, "hidden": false }, { "flags": [ "--data-dir" ], "dest": "data_dir", "type": "Path", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 153, "hidden": false }, { "flags": [ "--dry-run" ], "dest": "dry_run", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 154, "hidden": false }, { "flags": [ "--skip-install" ], "dest": "skip_install", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 155, "hidden": false }, { "flags": [ "--install-kaggle" ], "dest": "install_kaggle", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 156, "hidden": false }, { "flags": [ "--download-data" ], "dest": "download_data", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 157, "hidden": false }, { "flags": [ "--convert-raw" ], "dest": "convert_raw", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "", "line": 158, "hidden": false }, { "flags": [ "--raw-dir" ], "dest": "raw_dir", "type": "Path", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 159, "hidden": false }, { "flags": [ "--shard-dir" ], "dest": "shard_dir", "type": "Path", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 160, "hidden": false }, { "flags": [ "--shard-size" ], "dest": "shard_size", "type": "int", "default": 128, "required": false, "choices": null, "action": "store", "help": "", "line": 161, "hidden": false } ], "description": "" } ] }, { "command": "axosim-evaluate-model", "source_module": "axobench_iteration", "sections": [ { "name": "axosim-evaluate-model", "options": [ { "flags": [ "checkpoint" ], "dest": "checkpoint", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 205, "hidden": false }, { "flags": [ "--baseline-checkpoint" ], "dest": "baseline_checkpoint", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 207, "hidden": false, "mutually_exclusive_group": "baseline" }, { "flags": [ "--baseline-cache" ], "dest": "baseline_cache", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 208, "hidden": false, "mutually_exclusive_group": "baseline" }, { "flags": [ "--save-cache" ], "dest": "save_cache", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 209, "hidden": false }, { "flags": [ "--dataset-root" ], "dest": "dataset_root", "type": "str", "default": null, "required": true, "choices": null, "action": "store", "help": "", "line": 210, "hidden": false }, { "flags": [ "--model-id" ], "dest": "model_id", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 211, "hidden": false }, { "flags": [ "--device" ], "dest": "device", "type": "str", "default": "auto", "required": false, "choices": null, "action": "store", "help": "", "line": 212, "hidden": false }, { "flags": [ "--batch-size" ], "dest": "batch_size", "type": "int", "default": 128, "required": false, "choices": null, "action": "store", "help": "", "line": 213, "hidden": false }, { "flags": [ "--max-samples" ], "dest": "max_samples", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 214, "hidden": false }, { "flags": [ "--max-intervention-samples" ], "dest": "max_intervention_samples", "type": "int", "default": null, "required": false, "choices": null, "action": "store", "help": "", "line": 215, "hidden": false }, { "flags": [ "--max-interventions" ], "dest": "max_interventions", "type": "int", "default": 8, "required": false, "choices": null, "action": "store", "help": "", "line": 216, "hidden": false }, { "flags": [ "--intervention-workers" ], "dest": "intervention_workers", "type": "int", "default": 4, "required": false, "choices": null, "action": "store", "help": "", "line": 217, "hidden": false }, { "flags": [ "--interventions" ], "dest": "interventions", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "comma-separated intervention names; default discovers all", "line": 218, "hidden": false }, { "flags": [ "--multiple-morphologies" ], "dest": "multiple_morphologies", "type": "flag", "default": false, "required": false, "choices": null, "action": "store_true", "help": "balance the 120/40 evaluation profile across available morphologies", "line": 222, "hidden": false }, { "flags": [ "--morphology-id" ], "dest": "morphology_id", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "specific morphology for the default single-morphology evaluation", "line": 227, "hidden": false }, { "flags": [ "--ignore-start" ], "dest": "ignore_start", "type": "int", "default": 500, "required": false, "choices": null, "action": "store", "help": "", "line": 231, "hidden": false }, { "flags": [ "--calibration-fraction" ], "dest": "calibration_fraction", "type": "float", "default": 0.25, "required": false, "choices": null, "action": "store", "help": "", "line": 232, "hidden": false }, { "flags": [ "--default-spike-threshold" ], "dest": "default_spike_threshold", "type": "float", "default": 0.0, "required": false, "choices": null, "action": "store", "help": "", "line": 233, "hidden": false }, { "flags": [ "--bootstrap-replicates" ], "dest": "bootstrap_replicates", "type": "int", "default": 250, "required": false, "choices": null, "action": "store", "help": "", "line": 234, "hidden": false }, { "flags": [ "--bootstrap-seed" ], "dest": "bootstrap_seed", "type": "int", "default": 20260711, "required": false, "choices": null, "action": "store", "help": "", "line": 235, "hidden": false }, { "flags": [ "--output" ], "dest": "output", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "optional .json or .csv report path", "line": 237, "hidden": false, "mutually_exclusive_group": "output" }, { "flags": [ "--output-json" ], "dest": "output_json", "type": "str", "default": null, "required": false, "choices": null, "action": "store", "help": "argparse.SUPPRESS", "line": 238, "hidden": true, "mutually_exclusive_group": "output" } ], "description": "" } ] } ] }