def pred_step( self, x: tp.Any, states: types.States, initializing: bool, training: bool, ) -> PredStep: raise types.MissingMethod()
def init_step( self, x: tp.Any, y_true: tp.Any, sample_weight: tp.Optional[np.ndarray], class_weight: tp.Optional[np.ndarray], states: types.States, ) -> types.States: raise types.MissingMethod()
def train_step( self, x: tp.Any, y_true: tp.Any, sample_weight: tp.Optional[np.ndarray], class_weight: tp.Optional[np.ndarray], states: types.States, initializing: bool, training: bool, ) -> TrainStep: raise types.MissingMethod()
def summary_step( self, x: tp.Any, states: types.States, ) -> tp.List[types.SummaryTableEntry]: raise types.MissingMethod()
def train_step(self, *args, **kwargs): raise types.MissingMethod()
def reset_metrics(self) -> None: raise types.MissingMethod()
def train_step( self: M, inputs: tp.Any, labels: tp.Mapping[str, tp.Any], ) -> TrainStepOutput[M]: raise types.MissingMethod()
def grad_step( self: M, inputs: tp.Any, labels: tp.Mapping[str, tp.Any], ) -> GradStepOutput[M]: raise types.MissingMethod()
def pred_step( self: M, inputs: tp.Any, ) -> PredStepOutput[M]: raise types.MissingMethod()
def init_step( self: M, key: jnp.ndarray, inputs: tp.Any, ) -> M: raise types.MissingMethod()