jaxsnn.event.training

Classes

Any(*args, **kwargs)	Special type indicating an unconstrained type.
OptState(opt_state, params, …)	Container for the optimizer/training state across steps.

Functions

jaxsnn.event.training.data_loader(dataset: Tuple[Any, Any], batch_size: int, num_batches: Optional[int, None] = None, rng: Optional[jax.Array, None] = None)

jaxsnn.event.training.epoch(update_fn: Callable, test_fn: Callable[[List[jax.Array], Tuple[jax.Array, jax.Array]], Tuple[Any, str]], trainset, testset, batch_size_train: int, batch_size_test: int, state: jaxsnn.event.types.OptState, i: int)

jaxsnn.event.training.time_it(timed_function: Callable, *args) → Tuple[Any, float]

jaxsnn.event.training.update(optimizer, loss_fn: Callable, state: jaxsnn.event.types.OptState, batch: Tuple[Dict[str, jaxsnn.event.types.Spike], jax.Array]) → Tuple[jaxsnn.event.types.OptState, Tuple[jax.Array, Dict[str, jax.Array]]]