# Path: jax_verify/extensions/functional_lagrangian/lagrangian_form.py
# def _flatten_spatial_dims(x: Tensor) -> Tensor:
# def size_from_shape(shape: Shape) -> int:
#   def __init__(self, name):
#   def _init_params_per_sample(self, key: PRNGKey, *args) -> Params:
#   def init_params(self, key, *args, **kwargs):
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def process_params(self, lagrange_params: Params):
#   def name(self):
#   def __init__(self):
#   def _init_params_per_sample(self,
#                               key: PRNGKey,
#                               l_shape: Shape,
#                               init_zeros: bool = True) -> Params:
#   def _apply_per_sample(self, x: Tensor, lagrange_params: Params,
#                         step: int) -> Tensor:
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def __init__(self):
#   def _init_params_per_sample(self,
#                               key: PRNGKey,
#                               l_shape: Shape,
#                               init_zeros: bool = False) -> Params:
#   def _apply_per_sample(self, x: Tensor, lagrange_params: Params,
#                         step: int) -> Tensor:
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
# def get_lagrangian_form(config_lagrangian_form: ConfigDict) -> LagrangianForm:
# class LagrangianForm(metaclass=abc.ABCMeta):
# class Linear(LagrangianForm):
# class LinearExp(LagrangianForm):


from absl.testing import absltest
from absl.testing import parameterized
from jax_verify.extensions.functional_lagrangian import lagrangian_form
import chex
import haiku as hk
import jax
# 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.

"""Elementary tests for the Lagrangian forms."""


INPUT_SHAPES = (('batched_0d', [1]), ('batched_1d', [1, 2]),
                ('batched_2d', [1, 2, 3]), ('batched_3d', [1, 2, 3, 4]))


class ShapeTest(chex.TestCase):

  def setUp(self):
    super(ShapeTest, self).setUp()
    self._prng_seq = hk.PRNGSequence(13579)

  def _assert_output_shape(self, form, shape):
    x = jax.random.normal(next(self._prng_seq), shape)
    params = form.init_params(next(self._prng_seq), x.shape[1:])
    out = form.apply(x, params, step=0)
    assert out.ndim == 1

  @parameterized.named_parameters(*INPUT_SHAPES)
  def test_linear(self, shape):
    form = lagrangian_form.Linear()
    self._assert_output_shape(form, shape)

  @parameterized.named_parameters(*INPUT_SHAPES)
  def test_linear_exp(self, shape):

# Path: jax_verify/extensions/functional_lagrangian/lagrangian_form.py
# def _flatten_spatial_dims(x: Tensor) -> Tensor:
# def size_from_shape(shape: Shape) -> int:
#   def __init__(self, name):
#   def _init_params_per_sample(self, key: PRNGKey, *args) -> Params:
#   def init_params(self, key, *args, **kwargs):
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def process_params(self, lagrange_params: Params):
#   def name(self):
#   def __init__(self):
#   def _init_params_per_sample(self,
#                               key: PRNGKey,
#                               l_shape: Shape,
#                               init_zeros: bool = True) -> Params:
#   def _apply_per_sample(self, x: Tensor, lagrange_params: Params,
#                         step: int) -> Tensor:
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
#   def __init__(self):
#   def _init_params_per_sample(self,
#                               key: PRNGKey,
#                               l_shape: Shape,
#                               init_zeros: bool = False) -> Params:
#   def _apply_per_sample(self, x: Tensor, lagrange_params: Params,
#                         step: int) -> Tensor:
#   def _apply(self, x: Tensor, lagrange_params: Params, step: int) -> Tensor:
# def get_lagrangian_form(config_lagrangian_form: ConfigDict) -> LagrangianForm:
# class LagrangianForm(metaclass=abc.ABCMeta):
# class Linear(LagrangianForm):
# class LinearExp(LagrangianForm):


# Path: jax_verify/tests/functional_lagrangian/lagrangian_form_test.py
from absl.testing import absltest
from absl.testing import parameterized
from jax_verify.extensions.functional_lagrangian import lagrangian_form
import chex
import haiku as hk
import jax

# 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.

"""Elementary tests for the Lagrangian forms."""


INPUT_SHAPES = (('batched_0d', [1]), ('batched_1d', [1, 2]),
                ('batched_2d', [1, 2, 3]), ('batched_3d', [1, 2, 3, 4]))


class ShapeTest(chex.TestCase):

  def setUp(self):
    super(ShapeTest, self).setUp()
    self._prng_seq = hk.PRNGSequence(13579)

  def _assert_output_shape(self, form, shape):
    x = jax.random.normal(next(self._prng_seq), shape)
    params = form.init_params(next(self._prng_seq), x.shape[1:])
    out = form.apply(x, params, step=0)
    assert out.ndim == 1

  @parameterized.named_parameters(*INPUT_SHAPES)
  def test_linear(self, shape):
    form = lagrangian_form.Linear()
    self._assert_output_shape(form, shape)

  @parameterized.named_parameters(*INPUT_SHAPES)
  def test_linear_exp(self, shape):

    form = lagrangian_form.LinearExp()