# Path: fedjax/core/regularizers.py
# def _l2_regularize(params: Params, weight: float,
#                    center_params: Optional[Params],
#                    params_weights: Optional[Params]) -> float:
# def l2_regularizer(
#     weight: float,
#     center_params: Optional[Params] = None,
#     params_weights: Optional[Params] = None) -> Callable[[Params], float]:
#   def func(params):


from absl.testing import absltest
from fedjax.core import regularizers
import jax
import jax.numpy as jnp
# 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.
"""Tests for fedjax.core.regularizers."""





class RegularizersTest(absltest.TestCase):

  def test_l2_regularizer(self):
    params = {
        'linear0': {
            'w': jnp.array([1., 2., 3.])
        },
        'linear1': {
            'w': jnp.array([4., 5., 6.]),
            'b': jnp.array([1., 1., 1.])
        }
    }

    with self.subTest('weight'):
      regularizer = regularizers.l2_regularizer(weight=0.2)
      self.assertAlmostEqual(regularizer(params), 0.2 * 94, places=5)

    with self.subTest('center_params'):
      center_params = jax.tree_util.tree_map(jnp.ones_like, params)

# Path: fedjax/core/regularizers.py
# def _l2_regularize(params: Params, weight: float,
#                    center_params: Optional[Params],
#                    params_weights: Optional[Params]) -> float:
# def l2_regularizer(
#     weight: float,
#     center_params: Optional[Params] = None,
#     params_weights: Optional[Params] = None) -> Callable[[Params], float]:
#   def func(params):


# Path: fedjax/core/regularizers_test.py
from absl.testing import absltest
from fedjax.core import regularizers
import jax
import jax.numpy as jnp

# 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.
"""Tests for fedjax.core.regularizers."""





class RegularizersTest(absltest.TestCase):

  def test_l2_regularizer(self):
    params = {
        'linear0': {
            'w': jnp.array([1., 2., 3.])
        },
        'linear1': {
            'w': jnp.array([4., 5., 6.]),
            'b': jnp.array([1., 1., 1.])
        }
    }

    with self.subTest('weight'):
      regularizer = regularizers.l2_regularizer(weight=0.2)
      self.assertAlmostEqual(regularizer(params), 0.2 * 94, places=5)

    with self.subTest('center_params'):
      center_params = jax.tree_util.tree_map(jnp.ones_like, params)

      regularizer = regularizers.l2_regularizer(