Skip to main content

partitioned_module_parameter_group

eredu_runtime::parallel

Function partitioned_module_parameter_group 

Source 
pub fn partitioned_module_parameter_group<T, M>(
    logical_name: impl Into<String>,
    role: ParameterRole,
    preferred_units: usize,
    module: &M,
    sharding: impl FnMut(&ParameterMetadata, &[usize]) -> Result<MemberSharding, ParallelPlanError>,
) -> Result<ParameterGroupSpec, ParallelPlanError>
where
    T: Tensor,
    M: Parameterized<T>,
Expand description

Describes every parameter in a neutral module as one shared logical partition.