# Path: rigl/experimental/jax/training/training.py
# LABELKEY = dataset_base.ImageDataset.LABELKEY
# DATAKEY = dataset_base.ImageDataset.DATAKEY
# def _shard_batch(xs):
#   def _prepare(x):
# def train_step(
#     optimizer: flax.optim.Optimizer, batch: Mapping[str, jnp.array],
#     rng: Callable[[int], jnp.array], state: flax.deprecated.nn.Collection,
#     learning_rate_fn: Callable[[int], float]
# ) -> Tuple[flax.optim.Optimizer, flax.deprecated.nn.Collection, float, float]:
#   def loss_fn(
#       model: flax.deprecated.nn.Model
#   ) -> Tuple[float, Tuple[flax.deprecated.nn.Collection, jnp.array]]:
#   def __init__(
#       self,
#       optimizer_def: flax.optim.OptimizerDef,
#       initial_model: flax.deprecated.nn.Model,
#       initial_state: flax.deprecated.nn.Collection,
#       dataset: jnp.array,
#       rng: Callable[[int], jnp.array] = None,
#       summary_writer: Optional[tf.summary.SummaryWriter] = None,
#   ):
# class Trainer:
# 
# Path: rigl/experimental/jax/utils/utils.py
# def cross_entropy_loss(log_softmax_logits,
#                        labels):
# def compute_metrics(logits,
#                     labels):
# def _np_converter(obj):
# def dump_dict_json(data_dict, path):
# def count_param(model,
#                 param_names):
# def cosine_similarity(a, b):
# def param_as_array(params):
# def cosine_similarity_model(initial_model,
#                             current_model):
# def vector_difference_norm_model(initial_model,
#                                  current_model):
# def pairwise_longest(iterable):
# T = TypeVar('T')


import functools
import json
import operator
import tempfile
import flax
import jax
import jax.numpy as jnp
import numpy as np
from typing import Optional, Sequence, TypeVar
from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.training import training
from rigl.experimental.jax.utils import utils

    p_compute_metrics = jax.pmap(utils.compute_metrics, axis_name='batch')
    metrics = p_compute_metrics(logits, labels_correct)
    loss = metrics['loss']
    accuracy = metrics['accuracy']

    with self.subTest(name='loss_type'):
      self.assertIsInstance(loss, jnp.ndarray)

    with self.subTest(name='loss_len'):
      self.assertEqual(loss.size, 1)

    with self.subTest(name='loss_values'):
      self.assertGreaterEqual(loss.all(), 0)

    with self.subTest(name='accuracy_type'):
      self.assertIsInstance(accuracy, jnp.ndarray)

    with self.subTest(name='accuracy_Len'):
      self.assertEqual(accuracy.size, 1)

    with self.subTest(name='accuracy_values'):
      self.assertAlmostEqual(accuracy.all(), 1.0)

  def test_compute_metrics_incorrect(self):
    """Tests output when logit outputs indicate incorrect classification."""
    logits, labels_incorrect = self._create_logits_labels(False)
    logits = training._shard_batch(logits)
    labels_incorrect = training._shard_batch(labels_incorrect)


# Path: rigl/experimental/jax/training/training.py
# LABELKEY = dataset_base.ImageDataset.LABELKEY
# DATAKEY = dataset_base.ImageDataset.DATAKEY
# def _shard_batch(xs):
#   def _prepare(x):
# def train_step(
#     optimizer: flax.optim.Optimizer, batch: Mapping[str, jnp.array],
#     rng: Callable[[int], jnp.array], state: flax.deprecated.nn.Collection,
#     learning_rate_fn: Callable[[int], float]
# ) -> Tuple[flax.optim.Optimizer, flax.deprecated.nn.Collection, float, float]:
#   def loss_fn(
#       model: flax.deprecated.nn.Model
#   ) -> Tuple[float, Tuple[flax.deprecated.nn.Collection, jnp.array]]:
#   def __init__(
#       self,
#       optimizer_def: flax.optim.OptimizerDef,
#       initial_model: flax.deprecated.nn.Model,
#       initial_state: flax.deprecated.nn.Collection,
#       dataset: jnp.array,
#       rng: Callable[[int], jnp.array] = None,
#       summary_writer: Optional[tf.summary.SummaryWriter] = None,
#   ):
# class Trainer:
# 
# Path: rigl/experimental/jax/utils/utils.py
# def cross_entropy_loss(log_softmax_logits,
#                        labels):
# def compute_metrics(logits,
#                     labels):
# def _np_converter(obj):
# def dump_dict_json(data_dict, path):
# def count_param(model,
#                 param_names):
# def cosine_similarity(a, b):
# def param_as_array(params):
# def cosine_similarity_model(initial_model,
#                             current_model):
# def vector_difference_norm_model(initial_model,
#                                  current_model):
# def pairwise_longest(iterable):
# T = TypeVar('T')


# Path: rigl/experimental/jax/utils/utils_test.py
import functools
import json
import operator
import tempfile
import flax
import jax
import jax.numpy as jnp
import numpy as np
from typing import Optional, Sequence, TypeVar
from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.training import training
from rigl.experimental.jax.utils import utils


    p_compute_metrics = jax.pmap(utils.compute_metrics, axis_name='batch')
    metrics = p_compute_metrics(logits, labels_correct)
    loss = metrics['loss']
    accuracy = metrics['accuracy']

    with self.subTest(name='loss_type'):
      self.assertIsInstance(loss, jnp.ndarray)

    with self.subTest(name='loss_len'):
      self.assertEqual(loss.size, 1)

    with self.subTest(name='loss_values'):
      self.assertGreaterEqual(loss.all(), 0)

    with self.subTest(name='accuracy_type'):
      self.assertIsInstance(accuracy, jnp.ndarray)

    with self.subTest(name='accuracy_Len'):
      self.assertEqual(accuracy.size, 1)

    with self.subTest(name='accuracy_values'):
      self.assertAlmostEqual(accuracy.all(), 1.0)

  def test_compute_metrics_incorrect(self):
    """Tests output when logit outputs indicate incorrect classification."""
    logits, labels_incorrect = self._create_logits_labels(False)
    logits = training._shard_batch(logits)
    labels_incorrect = training._shard_batch(labels_incorrect)


    p_compute_metrics = jax.pmap(utils.compute_metrics, axis_name='batch')