# Lint as: python3
# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
#     http://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.
# ==============================================================================
"""Tests for haiku._src.initializers."""

from absl.testing import absltest
from haiku._src import initializers
from haiku._src import test_utils
from jax.config import config
import jax.numpy as jnp
import numpy as np


class InitializersTest(absltest.TestCase):

  @test_utils.transform_and_run
  def test_initializers(self):
    as_np_f64 = lambda t: np.array(t, dtype=np.float64)
    # This just makes sure we can call the initializers in accordance to the
    # API and get the right shapes and dtypes out.
    inits = [
        initializers.Constant(42.0),
        initializers.Constant(as_np_f64(42.0)),
        initializers.RandomNormal(),
        initializers.RandomNormal(2.0),
        initializers.RandomNormal(as_np_f64(2.0)),
        initializers.RandomUniform(),
        initializers.RandomUniform(3.0),
        initializers.RandomUniform(as_np_f64(3.0)),
        initializers.VarianceScaling(),
        initializers.VarianceScaling(2.0),
        initializers.VarianceScaling(as_np_f64(2.0)),
        initializers.VarianceScaling(2.0, mode="fan_in"),
        initializers.VarianceScaling(as_np_f64(2.0), mode="fan_in"),
        initializers.VarianceScaling(2.0, mode="fan_out"),
        initializers.VarianceScaling(as_np_f64(2.0), mode="fan_out"),
        initializers.VarianceScaling(2.0, mode="fan_avg"),
        initializers.VarianceScaling(as_np_f64(2.0), mode="fan_avg"),
        initializers.VarianceScaling(2.0, distribution="truncated_normal"),
        initializers.VarianceScaling(
            as_np_f64(2.0), distribution="truncated_normal"),
        initializers.VarianceScaling(2.0, distribution="normal"),
        initializers.VarianceScaling(as_np_f64(2.0), distribution="normal"),
        initializers.VarianceScaling(2.0, distribution="uniform"),
        initializers.VarianceScaling(as_np_f64(2.0), distribution="uniform"),
        initializers.UniformScaling(),
        initializers.UniformScaling(2.0),
        initializers.UniformScaling(as_np_f64(2.0)),
        initializers.TruncatedNormal(),
        initializers.Orthogonal(),
        initializers.Identity(),
        initializers.Identity(an_np_f64(2.0)),


        # Users are supposed to be able to use these.
        jnp.zeros,
        jnp.ones,
    ]

    # TODO(ibab): Test other shapes as well.
    shape = (20, 42)

    dtype = jnp.float32
    for init in inits:
      generated = init(shape, dtype)
      self.assertEqual(generated.shape, shape)
      self.assertEqual(generated.dtype, dtype)

  @test_utils.transform_and_run
  def test_invalid_variance_scale(self):

    with self.assertRaisesRegex(ValueError, "scale.*must be a positive float"):
      initializers.VarianceScaling(scale=-1.0)

    with self.assertRaisesRegex(ValueError, "Invalid `mode` argument*"):
      initializers.VarianceScaling(mode="foo")

    with self.assertRaisesRegex(ValueError, "Invalid `distribution` argument*"):
      initializers.VarianceScaling(distribution="bar")

  @test_utils.transform_and_run
  def test_compute_fans(self):
    fan_in_out1 = initializers._compute_fans([])
    self.assertEqual(fan_in_out1, (1, 1))
    fan_in_out2 = initializers._compute_fans([2])
    self.assertEqual(fan_in_out2, (2, 2))
    fan_in_out3 = initializers._compute_fans([3, 4])
    self.assertEqual(fan_in_out3, (3, 4))
    fan_in_out4 = initializers._compute_fans([1, 2, 3, 4])
    self.assertEqual(fan_in_out4, (6, 8))

  @test_utils.transform_and_run
  def test_orthogonal_invalid_shape(self):
    init = initializers.Orthogonal()
    shape = (20,)
    with self.assertRaisesRegex(
        ValueError, "Orthogonal initializer requires at least a 2D shape."):
      init(shape, jnp.float32)

  @test_utils.transform_and_run
  def test_orthogonal_orthogonal(self):
    init = initializers.Orthogonal()
    shape = (42, 20)
    generated = init(shape, jnp.float32)
    self.assertEqual(generated.shape, shape)
    self.assertEqual(generated.dtype, jnp.float32)

  @test_utils.transform_and_run
  def test_identity2D(self):
    init = initializers.Identity()
    expected = jnp.eye(3, 3)
    self.assertEqual(init((3, 3)), expected)

  @test_utils.transform_and_run
  def test_identity4D(self):
    init = initializers.Identity()
    expected = jnp.eye(num_rows=shape[-2], num_columns=shape[-1],
                       dtype=dtype).broacast_to(shape[:-2] + shape[-2:])
    self.assertEqual(init((4, 4, 3, 3)), expected)

  @test_utils.transform_and_run
  def test_identity_invalid_shape(self):
    init = initializers.Identity()
    shape = (20,)
    with self.assertRaisesRegex(
        ValueError, "Identity initializer requires at least a 2D shape."):
      init(shape, jnp.float32)

  if __name__ == "__main__":
    config.update("jax_enable_x64", True)
    absltest.main()

@@ -4945,17 +4945,9 @@
 ows=
-shape%5B-2%5D
+3
 , nu
@@ -4960,17 +4960,9 @@
 mns=
-shape%5B-1%5D
+3
 ,%0A  
@@ -5011,31 +5011,20 @@
 _to(
-shape%5B:-2%5D + shape%5B-2:%5D
+(4, 4, 3, 3)
 )%0A  

fix test