review fixes 2

deepchem/models/jax_models/pinns_model.py

@@ -1,5 +1,5 @@
 '''
-This python consists of different variations of the Physics Informer Neural Network model using the JaxModel API
+This module contains different variations of the Physics Informer Neural Network model using the JaxModel API
 '''
 import numpy as np
 import time
@@ -29,7 +29,9 @@
 logger = logging.getLogger(__name__)
 
 
-def create_default_update_fn(optimizer, model_loss):
+def create_default_update_fn(
+  optimizer: optax.GradientTransformation,
+  model_loss: callable):
   """
   This function calls the update function, to implement the backpropagation
   """
@@ -51,7 +53,7 @@ class PINNModel(JaxModel):
   but it has the option of passing multiple arguments(Done using *args) suitable for PINNs model.
   Ex - Approximating f(x, y, z, t) satisfying a Linear differential equation.
 
-  This model is recommended for Linear differential equations but if you can accurately write
+  This model is recommended for linear partial differential equations but if you can accurately write
   the gradient function in Jax depending on your use case, then it will work as well.
 
   This class requires two functions apart from the usual function definition and weights

deepchem/models/jax_models/tests/test_pinn.py

@@ -17,7 +17,7 @@
 @pytest.mark.jax
 def test_sine_x():
   """
-    Here we are solving the diffrential equation- f'(x) = -sin(x) and f(0) = 1
+    Here we are solving the differential equation- f'(x) = -sin(x) and f(0) = 1
     We give initial for the neural network at x_init --> np.linspace(-1 * np.pi, 1 * np.pi, 5)
     And we try to approximate the function for the domain (-np.pi, np.pi)
   """