diff --git a/gnm/shape/gnm_jax.py b/gnm/shape/gnm_jax.py
index 88c589c3..c03abe8d 100644
--- a/gnm/shape/gnm_jax.py
+++ b/gnm/shape/gnm_jax.py
@@ -16,21 +16,23 @@
 
 Usage:
   ```
-  gnm = gnm_jax.from_local(version=GNMVersion.V3, variant=GNMVariant.HEAD)
+  gnm = gnm_jax.GNM.from_local(
+      version=gnm_jax.GNMMajorVersion.V3, variant=gnm_jax.GNMVariant.HEAD
+  )
 
   # Generate batches of parameters.
   n_batch = 5
-  identity = np.random.uniform(shape=(n_batch, gnm.identity_dim))
-  expression = np.random.uniform(shape=(n_batch, gnm.expression_dim))
-  rotations = np.random.uniform(shape=[n_batch, gnm.num_joints, 3])
-  translation = np.random.uniform(shape=(n_batch, 3))
+  identity = np.random.uniform(size=(n_batch, gnm.identity_dim))
+  expression = np.random.uniform(size=(n_batch, gnm.expression_dim))
+  rotations = np.random.uniform(size=[n_batch, gnm.num_joints, 3])
+  translation = np.random.uniform(size=(n_batch, 3))
   vertices = gnm(identity, expression, rotations, translation)
   ```
 
   # This module is differentiable.
   grad_func = jax.grad(
      lambda *args: jnp.square(gnm(*args)).mean(),
-     argnums=np.Array([0, 1, 2, 3]))
+     argnums=(0, 1, 2, 3))
   grads = grad_func(identity, expression, rotations, translation)
 """
 
diff --git a/gnm/shape/gnm_numpy.py b/gnm/shape/gnm_numpy.py
index 37340d12..fe8bb15a 100644
--- a/gnm/shape/gnm_numpy.py
+++ b/gnm/shape/gnm_numpy.py
@@ -12,11 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""NumPy implementation of the GNM model .
+"""NumPy implementation of the GNM model.
 
 Example usage:
   ```
-  gnm = gnm_numpy.from_local(version=GNMVersion.V3, variant=GNMVariant.HEAD)
+  gnm = gnm_numpy.GNM.from_local(
+      version=gnm_numpy.GNMMajorVersion.V3, variant=gnm_numpy.GNMVariant.HEAD
+  )
 
   # Generate random identity, expression, rotations, translation parameters
   identity = np.random.normal(size=gnm.identity_dim)
diff --git a/gnm/shape/gnm_pytorch.py b/gnm/shape/gnm_pytorch.py
index 8757a263..76cde071 100644
--- a/gnm/shape/gnm_pytorch.py
+++ b/gnm/shape/gnm_pytorch.py
@@ -16,7 +16,10 @@
 
 Usage:
   ```
-  gnm = gnm_pytorch.from_local(version=GNMVersion.V3, variant=GNMVariant.HEAD)
+  gnm = gnm_pytorch.GNM.from_local(
+      version=gnm_pytorch.GNMMajorVersion.V3,
+      variant=gnm_pytorch.GNMVariant.HEAD,
+  )
 
   # Generate batches of parameters.
   n_batch = 5
diff --git a/gnm/shape/gnm_tensorflow.py b/gnm/shape/gnm_tensorflow.py
index 059f2dc4..ec7833f2 100644
--- a/gnm/shape/gnm_tensorflow.py
+++ b/gnm/shape/gnm_tensorflow.py
@@ -16,9 +16,9 @@
 
 Usage:
   ```
-  gnm = gnm_tensorflow.from_local(
-      version=GNMVersion.V3,
-      variant=GNMVariant.HEAD
+  gnm = gnm_tensorflow.GNM.from_local(
+      version=gnm_tensorflow.GNMMajorVersion.V3,
+      variant=gnm_tensorflow.GNMVariant.HEAD,
   )
 
   # Generate batches of parameters.