code: def test_tied_auto_encoder(self):
      ae = TiedAutoEncoder(latents=2, features=4)
      x = jnp.ones((1, ae.features))
      x_r, variables = init(ae)(random.key(0), x)
  
      param_shapes = unfreeze(
          jax.tree_util.tree_map(jnp.shape, variables["params"])
      )
      self.assertEqual(
          param_shapes,
          {
              "kernel": (4, 2),
          },
      )
      self.assertEqual(x.shape, x_r.shape)
  
file_path: flax/tests/core/design/core_tied_autoencoder_test.py
imports: from dataclasses import dataclass
  import jax
  from absl.testing import absltest
  from jax import numpy
  from jax import random
  from flax.core import init, lift, nn, unfreeze
label_desc: Use image node_id 1 for calling the TiedAutoEncoderTest obj's underlying member method code with example usage: obj.test_tied_auto_encoder() without return types