class_docstr: Tests for modules in model_utils.py.
class_docstr_tok: ["Tests","for","modules","in","model_utils.py","."]
docstr: Tests fo get_axial_1d_input and get_axial_2d_input.
docstr_tok: ["Tests","fo","get_axial_1d_input","and","get_axial_2d_input","."]
code: def test_axial_reshaping_utils(self, axis):
      """Tests fo get_axial_1d_input and get_axial_2d_input."""
      input_shape = (4, 8, 16, 32)  # Shape = `[bs, h, w, c]`
      inputs_2d = jnp.array(np.random.normal(size=input_shape))
      inputs_1d = model_utils.get_axial_1d_input(inputs_2d, axis=axis)
      inputs_back_to_2d = model_utils.get_axial_2d_input(
          inputs_1d, axis=axis, two_d_shape=input_shape
      )
      self.assertTrue(jnp.array_equal(inputs_2d, inputs_back_to_2d))
  
file_path: scenic/scenic/projects/fast_vit/tests/test_model_utils.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.projects.fast_vit import model_utils
label_desc: Use image node_id 5 for calling the AttentionLayersTest obj's underlying member method code with example usage: obj.test_axial_reshaping_utils(axis) without return types