kkthn.backbone

kkthn.backbone.glorot_init(key, in_dim: int, out_dim: int, *, dtype=<class 'jax.numpy.float64'>)

kkthn.backbone.init_mlp_params(key, layer_sizes: list[int] | tuple[int, ...], *, dtype=<class 'jax.numpy.float64'>)

kkthn.backbone.mlp_apply(params, x)

kkthn.backbone.make_batched_mlp_apply()