verified fit method

# JAX depend

import jax.numpy as jnp

import jax

from jax import random

import haiku as hk

import optax

class JaxModel(Model):

  """This is a DeepChem model implemented by a Jax Model

  Here is a simple example of that uses JaxModel to train a

  Haiku (JAX Nueral Network Library) based model on deepchem

  dataset.

  >> def f(x):

  >>   net = hk.nets.MLP([512, 256, 128, 1])

  >>   return net(x)

  >> model = hk.without_apply_rng(hk.transform(f))

  >> rng = jax.random.PRNGKey(500)

  >> x, _, _, _ = next(iter(train_dataset.iterbatches(batch_size=256)))

  >> params = model.init(rng, x)

  >> j_m = JaxModel(model, params, 256, 0.001, 100)

  >> j_m.fit(train_dataset)

  """

  def __init__(self,

    pass model.apply for model here

    """

    self.loss = lambda pred, tar: jnp.mean(optax.l2_loss(pred, tar))

    self.loss = loss  #lambda pred, tar: jnp.mean(optax.l2_loss(pred, tar))

    self.batch_size = batch_size

    self.optimizer = optax.adam(learning_rate)

    self.optimizer = optimizer

    self.model = model  # this is a function, hk.apply

    self.params = params

    self._built = False

    averaged_batches = 0

    loss = self.loss

    grad_update = self._create_gradient_fn(self.loss, self.optimizer)

    grad_update = self._create_gradient_fn(self.loss, self.model,

                                           self.optimizer)

    params, opt_state = self._get_trainable_params()

    for batch in generator:

    self.params = params

    self.opt_state = opt_state

  def _create_gradient_fn(self, loss, optimizer, p=None):

  def _create_gradient_fn(self, loss, model, optimizer, p=None):

    """

    This function calls the update function, to implement the backpropogation

    """

    def model_loss(params, batch, target):

      predict = model.apply(params, batch)

      return loss(predict, target)

    @jax.jit

    def update(params, opt_state, batch, target):

      batch_loss, grads = jax.value_and_grad(loss)(params, batch, target)