class_docstr: Tests for utilities in nn_ops.py.
class_docstr_tok: ["Tests","for","utilities","in","nn_ops.py","."]
docstr: Tests compute_relative_positions.
  
  Args:
    spatial_axis: position axis passed to the compute_relative_positions.
    expected_output_shape: expected shape of the output.
docstr_tok: ["Tests","compute_relative_positions",".","Args",":","spatial_axis",":","position","axis","passed","to","the","compute_relative_positions",".","expected_output_shape",":","expected","shape","of","the","output","."]
code: def test_compute_relative_positions(
      self, spatial_axis, expected_output_shape
  ):
      """Tests compute_relative_positions.
  
      Args:
        spatial_axis: position axis passed to the compute_relative_positions.
        expected_output_shape: expected shape of the output.
      """
      query_spatial_shape = (3, 5)
      key_spatial_shape = (4, 6)
      relative_positions = nn_ops.compute_relative_positions(
          query_spatial_shape, key_spatial_shape, spatial_axis
      )
  
      # test output shape
      self.assertEqual(relative_positions.shape, expected_output_shape)
  
      # test maximum positional distances
      for dim_i, dim in enumerate(spatial_axis):
          max_positional_distances = (
              query_spatial_shape[dim] + key_spatial_shape[dim] - 2
          )
          self.assertEqual(
              max_positional_distances,
              jnp.max(relative_positions[dim_i]),
          )
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import flax.linen
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_ops
label_desc: Use image node_id 1 for calling the NNOpsTest obj's underlying member method code with example usage: obj.test_compute_relative_positions(spatial_axis, expected_output_shape) without return types