Added things

fx/vmap.py

@@ -19,8 +19,8 @@
 #
 # How is this feat accomplished? One observation is that to "batch" a model, it
 # suffices to batch each individual operation. In other words, given an
-# operation that works on the current shape, how do we make it work on another
-# batch dimension? This leads us to batching rules.
+# operation that works on the current shape, how do we make it work with an
+# additional batch dimension? This leads us to batching rules.
 #
 # Batching Rules
 # ---------------
@@ -194,4 +194,4 @@ def forward(self, a, b):
 # outer product computation. ((B, N), (M,)) -> (B, N, M)
 
 model = vmap(model, in_axes=(0, None), example_args=(x[0], y))
-print(model(x, y).shape)  # ((3, 5), (2,)) -> (3, 5, 2)
+print(model(x, y).shape)  # ((3, 5), (2,)) -> (3, 5, 2)