code: def test_add(self):
      layout_map = distribution_lib.LayoutMap(self.device_mesh)
      layout_map["dense/kernel"] = self.sharded_2d
      layout_map["dense/bias"] = self.sharded_1d
      # Test for adding list/tuple as shortcut for TensorLayout
      layout_map["conv/bias"] = ("model",)
  
      # Make there are two items in the map, and we access them via the
      # underlying container at layout_map._layout_map
      self.assertLen(layout_map, 3)
  
      kernel_layout = layout_map["dense/kernel"]
      self.assertEqual(kernel_layout.axes, (None, "model"))
      self.assertIs(kernel_layout.device_mesh, self.device_mesh)
  
      bias_layout = layout_map["dense/bias"]
      self.assertEqual(bias_layout.axes, ("model",))
      self.assertIs(bias_layout.device_mesh, self.device_mesh)
  
      conv_bias_layout = layout_map["conv/bias"]
      self.assertEqual(conv_bias_layout.axes, ("model",))
      self.assertIs(bias_layout.device_mesh, self.device_mesh)
  
      with self.assertRaisesRegex(
          ValueError, "dense/kernel already exist"
      ):
          layout_map["dense/kernel"] = self.sharded_2d
  
      with self.assertRaisesRegex(
          ValueError, "should be a TensorLayout"
      ):
          layout_map["conv.kernel"] = ["a", "b"]
  
imports: import os
  from unittest import mock
  import numpy
  import pytest
  import tensorflow
  from keras import backend
  from keras import testing
  from keras.backend import distribution_lib
  from keras.distribution import distribution_lib
label_desc: Use image node_id 2 for calling the LayoutMapTest obj's underlying member method code with example usage: obj.test_add() without return types