belated updates

eringrant · May 30, 2024 · 529903d · 529903d
1 parent 1c36bd3
commit 529903d
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 2 deletions.
diff --git a/nets/models/feedforward.py b/nets/models/feedforward.py
@@ -64,7 +64,7 @@ def __init__(
 class MLP(eqx.Module):
   """Multi-layer perceptron."""
 
-  layers: tuple[Linear, ...]
+  layers: tuple[enn.Linear, ...]
   dropouts: tuple[enn.Dropout, ...]
   activation: Callable
   final_activation: Callable

diff --git a/nets/simulators/online_sgd.py b/nets/simulators/online_sgd.py
@@ -234,7 +234,7 @@ def simulate(
   exemplar_noise_scale: float,
   # Sampler params.
   sampler_cls: type[samplers.EpochSampler],  # TODO(eringrant): Use `SingletonSampler`.
-) -> tuple[pd.DataFrame, ...]:
+) -> tuple[eqx.Module, pd.DataFrame]:
   """Simulate in-context learning of classification tasks."""
   logging.info(f"Using JAX backend: {jax.default_backend()}\n")
 

diff --git a/py.typed b/py.typed
diff --git a/setup.cfg b/setup.cfg
@@ -30,5 +30,8 @@ install_requires =
     equinox>=0.11.0
     optax>=0.1.5
 
+[options.package_data]
+nets = py.typed
+
 [options.packages.find]
 include = nets*