<commit_before># Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from absl.testing import absltest
from absl.testing import parameterized

import jax
from jax import random
from jax import tree_util
import jax.numpy as jnp

from jaxopt import objective
from jaxopt import ScipyBoundedLeastSquares
from jaxopt import ScipyBoundedMinimize
from jaxopt import ScipyLeastSquares
from jaxopt import ScipyMinimize
from jaxopt import ScipyRootFinding
from jaxopt._src import scipy_wrappers
from jaxopt._src import test_util
from jaxopt._src.tree_util import tree_scalar_mul

import numpy as onp
import scipy as osp
from sklearn import datasets


class JnpToOnpTest(test_util.JaxoptTestCase):

  def setUp(self):
    super().setUp()

    m, n = 10, 9
    self.x_split_indices = [4, 6, 9]
    self.x_reshape_shape = [-1, 2]
    self.y_split_indices = [6, 8]
    self.y_reshape_shape = [-1, 3]

    key = random.PRNGKey(0)
    key_x, key_w1, key_w2 = random.split(key, 3)
    self.x = random.normal(key_x, [m])
    self.w1 = random.normal(key_w1, [m, m])
    self.w2 = random.normal(key_w2, [m, n])

  @staticmethod
  def _do_pytree(array, split_indices, reshape_shape):
    """Creates a (hardcoded) PyTree out of flat 1D array for testing purposes."""
    pytree = jnp.split(array, split_indices)
    pytree[0] = pytree[0].reshape(reshape_shape)
    pytree[-1] = {'-1': pytree[-1]}
    return pytree

  @staticmethod
  def _undo_pytree(pytree):
    """Recovers flat 1D array from (hardcoded) PyTree for testing purposes."""
    pytree = pytree.copy()
    pytree[0] = pytree[0].reshape([-1])
    pytree[-1] = pytree[-1]['-1']
    return jnp.concatenate(pytree)

  def test_vals_and_jac(self):
    fun_flat = lambda x: jnp.dot(jnp.arctan(jnp.dot(x, self.w1)), self.w2)

    def fun(x_pytree):
      """Wraps fun_flat with (hardcoded) yTree input / output."""
      x = self._undo_pytree(x_pytree)
      y = fun_flat(x)
      return self._do_pytree(y, self.y_split_indices, self.y_reshape_shape)

    # Tests function output.
    x_pytree = self._do_pytree(self.x,
                               self.x_split_indices,
                               self.x_reshape_shape)
    y_pytree = fun(x_pytree)
    self.assertArraysAllClose(fun_flat(self.x), self._undo_pytree(y_pytree))

    # Tests jnp_to_onp.
    self.assertArraysAllClose(self._undo_pytree(y_pytree),
                              scipy_wrappers.jnp_to_onp(y_pytree))

    # Tests Jacobian.
    x_pytree_topology = scipy_wrappers.pytree_topology_from_example(x_pytree)
    y_pytree_topology = scipy_wrappers.pytree_topology_from_example(y_pytree)
    jac_jnp_to_onp = scipy_wrappers.make_jac_jnp_to_onp(x_pytree_topology,
                                                        y_pytree_topology)

    jac_flat = jax.jacrev(fun_flat)(self.x)
    jac_pytree = jax.jacrev(fun)(x_pytree)
    self.assertArraysAllClose(jac_flat, jac_jnp_to_onp(jac_pytree))


class ScipyMinimizeTest(test_util.JaxoptTestCase):

  def setUp(self):
    super().setUp()

    self.n_samples, self.n_features, self.n_classes = 50, 5, 3
    self.data = datasets.make_classification(n_samples=self.n_samples,
                                             n_features=self.n_features,
                                             n_classes=self.n_classes,
                                             n_informative=3,
                                             random_state=0)
    self.default_l2reg = float(self.n_samples)

    self.solver_kwargs = {'method': 'L-BFGS-B',
                          'tol': 1e-3,
                          'options': {'maxiter': 500}}

    def logreg_fun(params, *args, **kwargs):
      params = tree_util.tree_leaves(params)
      return objective.l2_multiclass_logreg(params[0], *args, **kwargs)
    self.logreg_fun = logreg_fun

    def logreg_with_intercept_fun(params, *args, **kwargs):
      return objective.l2_multiclass_logreg_with_intercept(
          (params['W'], params['b']), *args, **kwargs)
    self.logreg_with_intercept_fun = logreg_with_intercept_fun

  def test_logreg(self):
    lbfgs = ScipyMinimize(fun=self.logreg_fun,
                          **self.solver_kwargs)<commit_msg>Fix: test_logreg_with_intercept.<commit_after># Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from absl.testing import absltest
from absl.testing import parameterized

import jax
from jax import random
from jax import tree_util
import jax.numpy as jnp

from jaxopt import objective
from jaxopt import ScipyBoundedLeastSquares
from jaxopt import ScipyBoundedMinimize
from jaxopt import ScipyLeastSquares
from jaxopt import ScipyMinimize
from jaxopt import ScipyRootFinding
from jaxopt._src import scipy_wrappers
from jaxopt._src import test_util
from jaxopt._src.tree_util import tree_scalar_mul

import numpy as onp
import scipy as osp
from sklearn import datasets


class JnpToOnpTest(test_util.JaxoptTestCase):

  def setUp(self):
    super().setUp()

    m, n = 10, 9
    self.x_split_indices = [4, 6, 9]
    self.x_reshape_shape = [-1, 2]
    self.y_split_indices = [6, 8]
    self.y_reshape_shape = [-1, 3]

    key = random.PRNGKey(0)
    key_x, key_w1, key_w2 = random.split(key, 3)
    self.x = random.normal(key_x, [m])
    self.w1 = random.normal(key_w1, [m, m])
    self.w2 = random.normal(key_w2, [m, n])

  @staticmethod
  def _do_pytree(array, split_indices, reshape_shape):
    """Creates a (hardcoded) PyTree out of flat 1D array for testing purposes."""
    pytree = jnp.split(array, split_indices)
    pytree[0] = pytree[0].reshape(reshape_shape)
    pytree[-1] = {'-1': pytree[-1]}
    return pytree

  @staticmethod
  def _undo_pytree(pytree):
    """Recovers flat 1D array from (hardcoded) PyTree for testing purposes."""
    pytree = pytree.copy()
    pytree[0] = pytree[0].reshape([-1])
    pytree[-1] = pytree[-1]['-1']
    return jnp.concatenate(pytree)

  def test_vals_and_jac(self):
    fun_flat = lambda x: jnp.dot(jnp.arctan(jnp.dot(x, self.w1)), self.w2)

    def fun(x_pytree):
      """Wraps fun_flat with (hardcoded) yTree input / output."""
      x = self._undo_pytree(x_pytree)
      y = fun_flat(x)
      return self._do_pytree(y, self.y_split_indices, self.y_reshape_shape)

    # Tests function output.
    x_pytree = self._do_pytree(self.x,
                               self.x_split_indices,
                               self.x_reshape_shape)
    y_pytree = fun(x_pytree)
    self.assertArraysAllClose(fun_flat(self.x), self._undo_pytree(y_pytree))

    # Tests jnp_to_onp.
    self.assertArraysAllClose(self._undo_pytree(y_pytree),
                              scipy_wrappers.jnp_to_onp(y_pytree))

    # Tests Jacobian.
    x_pytree_topology = scipy_wrappers.pytree_topology_from_example(x_pytree)
    y_pytree_topology = scipy_wrappers.pytree_topology_from_example(y_pytree)
    jac_jnp_to_onp = scipy_wrappers.make_jac_jnp_to_onp(x_pytree_topology,
                                                        y_pytree_topology)

    jac_flat = jax.jacrev(fun_flat)(self.x)
    jac_pytree = jax.jacrev(fun)(x_pytree)
    self.assertArraysAllClose(jac_flat, jac_jnp_to_onp(jac_pytree))


class ScipyMinimizeTest(test_util.JaxoptTestCase):

  def setUp(self):
    super().setUp()

    self.n_samples, self.n_features, self.n_classes = 50, 5, 3
    self.data = datasets.make_classification(n_samples=self.n_samples,
                                             n_features=self.n_features,
                                             n_classes=self.n_classes,
                                             n_informative=3,
                                             random_state=0)
    self.default_l2reg = float(self.n_samples)

    self.solver_kwargs = {'method': 'L-BFGS-B',
                          'tol': 1e-5,
                          'options': {'maxiter': 500}}

    def logreg_fun(params, *args, **kwargs):
      params = tree_util.tree_leaves(params)
      return objective.l2_multiclass_logreg(params[0], *args, **kwargs)
    self.logreg_fun = logreg_fun

    def logreg_with_intercept_fun(params, *args, **kwargs):
      return objective.l2_multiclass_logreg_with_intercept(
          (params['W'], params['b']), *args, **kwargs)
    self.logreg_with_intercept_fun = logreg_with_intercept_fun

  def test_logreg(self):
    lbfgs = ScipyMinimize(fun=self.logreg_fun,
                          **self.solver_kwargs)
{'fraction_non_alphanumeric': 0.09685333333333333, 'fraction_numerical': 0.012373333333333333, 'mean_word_length': 2.931236897274633, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7433891', 'n_tokens_mistral': 3311, 'n_tokens_neox': 3251, 'n_words': 704}