code: def test_lsh_self_attention_masked_non_causal(self):
      # Test that when the input that is in the masked area changes the attention
      # for the un-masked outputs doesn't change, but the masked region does
      # change.
      with fastmath.use_backend(fastmath.Backend.JAX):
          layer = efficient_attention.LSHSelfAttention(
              n_heads=5,
              d_qk=7,
              d_v=17,
              causal=False,
              masked=True,
              chunk_len=8,
              n_chunks_before=1,
              n_chunks_after=0,
              n_hashes=2,
              n_buckets=4,
              use_reference_code=True,
              attention_dropout=0.0,
              mode="train",
          )
  
          batch = 5
          max_len = 32
          hidden = 8
  
          x = np.random.uniform(size=(batch, max_len, hidden))
          mask = np.ones((batch, max_len)).astype(bool)
          rngs = jax.random.randint(
              jax.random.PRNGKey(0),
              (batch,),
              minval=1,
              maxval=max_len - 1,
          )
  
          # Set some suffix of each mask[b] to 0.
          for i in range(batch):
              mask[i, rngs[i] :] = 0
  
          # Fix rngs and get the output for the LSH layer.
          def get_output(x, mask):
              xs = [x, mask]
              _, _ = layer.init(
                  shapes.signature(xs), jax.random.PRNGKey(0)
              )
              return layer(xs, rng=jax.random.PRNGKey(1))
  
          # Get the attention output for masked x.
          y = get_output(x, mask)
  
          # Change x, but only in the masked regions.
          for i in range(batch):
              x[i, rngs[i] :] = np.random.uniform(
                  size=(max_len - rngs[i], hidden)
              )
  
          y2 = get_output(x, mask)
  
          for i in range(batch):
              # y and y2 should be identical in the non-masked part.
              np.testing.assert_array_almost_equal(
                  y[i, : rngs[i]], y2[i, : rngs[i]], decimal=6
              )
  
              # In the masked out part, they should be different.
              self.assertGreater(
                  np.mean(np.abs(y[i, rngs[i] :] - y2[i, rngs[i] :])),
                  1e-5,
              )
  
imports: from absl.testing import parameterized
  import jax
  import numpy
  from tensorflow import test
  from trax import fastmath
  from trax import shapes
  from trax.fastmath import numpy
  from trax.layers.research import efficient_attention
label_desc: Use image node_id 14 for calling the EfficientAttentionTest obj's underlying member method code with example usage: obj.test_lsh_self_attention_masked_non_causal() and returns: layer