code: def test_keras_layer():
      if not is_backend_tested("tensorflow"):
          pytest.skip()
      else:
          import tensorflow as tf
          from tensorflow.keras.models import Sequential
          from tensorflow.keras.layers import (
              Conv2D as Conv2d,
              Dense as Linear,
              ReLU,
          )
          from einops.layers.keras import (
              Rearrange,
              Reduce,
              EinMix,
              keras_custom_objects,
          )
  
          def create_keras_model():
              return Sequential(
                  [
                      Conv2d(6, kernel_size=5, input_shape=[32, 32, 3]),
                      Reduce(
                          "b c (h h2) (w w2) -> b c h w",
                          "max",
                          h2=2,
                          w2=2,
                      ),
                      Conv2d(16, kernel_size=5),
                      Reduce(
                          "b c (h h2) (w w2) -> b c h w",
                          "max",
                          h2=2,
                          w2=2,
                      ),
                      Rearrange("b c h w -> b (c h w)"),
                      Linear(120),
                      ReLU(),
                      Linear(84),
                      ReLU(),
                      EinMix(
                          "b c1 -> (b c2)",
                          weight_shape="c1 c2",
                          bias_shape="c2",
                          c1=84,
                          c2=84,
                      ),
                      EinMix(
                          "(b c2) -> b c3",
                          weight_shape="c2 c3",
                          bias_shape="c3",
                          c2=84,
                          c3=84,
                      ),
                      Linear(10),
                  ]
              )
  
          model1 = create_keras_model()
          model2 = create_keras_model()
          input = numpy.random.normal(size=[10, 32, 32, 3]).astype(
              "float32"
          )
          assert not numpy.allclose(
              model1.predict_on_batch(input),
              model2.predict_on_batch(input),
          )
  
          # get some temp filename
          with tempfile.NamedTemporaryFile(mode="r+b") as f:
              tmp_filename = f.name
          # save arch + weights
          print("temp_path_keras1", tmp_filename)
          tf.keras.models.save_model(model1, tmp_filename)
          model3 = tf.keras.models.load_model(
              tmp_filename, custom_objects=keras_custom_objects
          )
          assert numpy.allclose(
              model1.predict_on_batch(input),
              model3.predict_on_batch(input),
          )
  
          # save arch as json
          model4 = tf.keras.models.model_from_json(
              model1.to_json(), custom_objects=keras_custom_objects
          )
          model1.save_weights(tmp_filename)
          model4.load_weights(tmp_filename)
          model2.load_weights(tmp_filename)
          assert numpy.allclose(
              model1.predict_on_batch(input),
              model4.predict_on_batch(input),
          )
          assert numpy.allclose(
              model1.predict_on_batch(input),
              model2.predict_on_batch(input),
          )
  
file_path: einops/tests/test_layers.py
imports: import pickle
  import tempfile
  from collections import namedtuple
  import numpy
  import pytest
  from einops import rearrange, reduce
  from .None import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS
label_desc: Use image node_id 8 for calling a global function with example usage: test_keras_layer() and returns: Sequential