# coding=utf-8
# Copyright 2022 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for Praxis streaming normalization layers."""

from absl.testing import absltest
from absl.testing import parameterized
import numpy as np
from praxis import base_layer
from praxis.layers import streaming
from praxis.layers.streaming import test_utils

instantiate = base_layer.instantiate


class StreamingNormalizationTest(test_utils.TestCase):

  @parameterized.named_parameters(
      ('Basic',),
      ('Group4Dim2', 1, 4, 2),
      ('Group1', 1, 1),
      ('Stride2', 2),
      ('Stride2Group1', 2, 1),
      ('Stride4', 4),
  )
  def test_stream(self, stride=1, num_groups=2, dim=4):
    seqlen = 8
    batch, input_dim = 2, dim
    np.random.seed(None)
    inputs = np.random.normal(
        0.1, 0.5, [batch, seqlen, 1, input_dim]).astype(np.float32)
    paddings = np.array([[1, 1, 1, 0, 0, 0, 1, 1]]).astype(np.float32)
    paddings = np.concatenate((paddings, paddings), axis=0)

    context_p = base_layer.JaxContext.HParams(do_eval=True)
    with base_layer.JaxContext.new_context(hparams=context_p):
      p = streaming.GroupNorm.HParams(
          name='stream_gn',
          dim=input_dim,
          num_groups=num_groups,
          cumulative=True,
          input_rank=4)

    self.assertEqual(p.cls.get_stride(p), 1)
    self.assertEqual(p.cls.get_right_context(p), 0)
    self._compare_stream_non_stream(
        inputs, paddings, p, p, stride)


if __name__ == '__main__':
  absltest.main()

@@ -798,16 +798,57 @@
 reaming%0A
+from praxis.layers import normalizations%0A
 from pra
@@ -1238,17 +1238,18 @@
 eqlen = 
-8
+16
 %0A    bat
@@ -1250,16 +1250,21 @@
    batch
+_size
 , input_
@@ -1280,33 +1280,8 @@
 dim%0A
-    np.random.seed(None)%0A
     
@@ -1331,16 +1331,21 @@
 , %5Bbatch
+_size
 , seqlen
@@ -1402,39 +1402,48 @@
  np.
-array(%5B%5B1, 1, 1, 0, 0, 0, 1, 1%5D
+random.randint(0, 2, %5Bbatch_size, seqlen
 %5D).a
@@ -1452,223 +1452,53 @@
 ype(
-np.
+'
 float32
-)
+')%0A
 %0A    p
-addings = np.concatenate((paddings, paddings), axis=0)%0A%0A    context_p = base_layer.JaxContext.HParams(do_eval=True)%0A    with base_layer.JaxContext.new_context(hparams=context_p):%0A      p = streaming
+_non_stream = normalizations
 .Gro
@@ -1525,16 +1525,18 @@
     
-  
 name='
+non_
 stre
@@ -1543,18 +1543,16 @@
 am_gn',%0A
-  
         
@@ -1574,18 +1574,16 @@
         
-  
 num_grou
@@ -1605,18 +1605,16 @@
         
-  
 cumulati
@@ -1622,18 +1622,16 @@
 e=True,%0A
-  
         
@@ -1645,16 +1645,150 @@
 ank=4)%0A%0A
+    # Streaming aware layer%0A    p_stream = streaming.GroupNorm.HParams(name='stream_gn')%0A    p_stream.copy_fields_from(p_non_stream)%0A%0A
     self
@@ -1793,32 +1793,39 @@
 lf.assertEqual(p
+_stream
 .cls.get_stride(
@@ -1825,16 +1825,23 @@
 stride(p
+_stream
 ), 1)%0A  
@@ -1860,16 +1860,23 @@
 tEqual(p
+_stream
 .cls.get
@@ -1891,16 +1891,23 @@
 ontext(p
+_stream
 ), 0)%0A  
@@ -1968,19 +1968,37 @@
 dings, p
-, p
+_non_stream, p_stream
 , stride

Enable streaming support for pos_embedding and norm in SubSampleConvProjection. Also create streaming layers in setup() instead of override (it will allow to to control supported layer types with multiple options for streaming).