code: def test_e2e_model_parallel_model(self):
      shape = (4, 2)
      axis_names = ["batch", "model"]
      device_mesh = distribution_lib.DeviceMesh(
          shape, axis_names, backend_dlib.list_devices()
      )
  
      layout_map = distribution_lib.LayoutMap(device_mesh)
      layout_map[".*dense.*kernel"] = distribution_lib.TensorLayout(
          [None, "model"]
      )
      layout_map[".*dense.*bias"] = distribution_lib.TensorLayout(
          ["model"]
      )
  
      distribution = distribution_lib.ModelParallel(
          device_mesh, layout_map, batch_dim_name="batch"
      )
      with distribution.scope():
          inputs = layers.Input(shape=[28, 28, 1])
          y = layers.Flatten()(inputs)
          y = layers.Dense(
              units=200, use_bias=False, activation="relu"
          )(y)
          y = layers.Dropout(0.4)(y)
          y = layers.Dense(units=10, activation="softmax")(y)
          model = models.Model(inputs=inputs, outputs=y)
  
      for weight in model.weights:
          if "kernel" in weight.name:
              self.assertEqual(
                  weight._value.sharding.spec, (None, "model")
              )
          elif "bias" in weight.name:
              self.assertEqual(weight._value.sharding.spec, ("model",))
          else:
              self.assertTrue(
                  weight._value.sharding.is_fully_replicated
              )
  
      inputs = np.random.normal(size=(32, 28, 28, 1))
      labels = np.random.normal(size=(32, 10))
  
      with distribution.scope():
          model.compile(loss="mse")
          model.fit(inputs, labels)
  
imports: import functools
  import os
  from unittest import mock
  import jax
  import numpy
  import pytest
  from keras import backend
  from keras import layers
  from keras import models
  from keras import testing
  from keras.backend import distribution_lib
  from keras.distribution import distribution_lib
label_desc: Use image node_id 15 for calling the JaxDistributionLibTest obj's underlying member method code with example usage: obj.test_e2e_model_parallel_model() without return types