code: def test_rotary_monotonicity(self):
      layer = rotary_pe.Rotate()
      batch_size = 1
      seq_len = 32
      d_model = 512
      shape = (batch_size, seq_len, d_model)
      q, k = np.ones(shape).astype(np.float32), np.ones(shape).astype(
          np.float32
      )
      q, k = layer(q), layer(k)
  
      self.assertEqual(q.dtype, np.float32)
      self.assertEqual(q.shape, shape)
  
      # Test monotonicity of the resulting dot_product for the two first tokens
      # in close proximity
      dot_product = np.einsum("bnd, bmd -> bnm", q, k)
  
      self.assertTrue(
          (dot_product[0, 0, :9] > dot_product[0, 0, 1:10]).all()
      )
      self.assertTrue(
          (dot_product[0, 1, 1:10] > dot_product[0, 1, 2:11]).all()
      )
  
imports: from absl.testing import absltest
  import numpy
  from trax.layers.research import rotary_positional_embedding
label_desc: Use image node_id 1 for calling the RelAttentionTest obj's underlying member method code with example usage: obj.test_rotary_monotonicity() without return types