def assertSerializable(self, model, batch_size=<NUM_LIT:1>):
batch_shape = [batch_size]<EOL>input_shape = batch_shape + model.input.shape[<NUM_LIT:1>:].as_list()<EOL>dtype = model.input.dtype.as_numpy_dtype()<EOL>model_file = self.create_tempfile()<EOL>model.save(model_file.full_path)<EOL>model_copy = tfk.models.load_model(model_file.full_path)<EOL>x = np.random.uniform(-<NUM_LIT>, <NUM_LIT>, input_shape).astype(dtype)<EOL>model_x_mean = self.evaluate(model(x).mean())<EOL>self.assertAllEqual(model_x_mean, self.evaluate(model_copy(x).mean()))<EOL>output_shape = model_x_mean.shape<EOL>y = np.random.uniform(<NUM_LIT:0.>, <NUM_LIT:1.>, output_shape).astype(dtype)<EOL>self.assertAllEqual(self.evaluate(model(x).log_prob(y)),<EOL>self.evaluate(model_copy(x).log_prob(y)))<EOL>
Assert that a model can be saved/loaded via Keras Model.save/load_model.

        Arguments:
          model: A Keras model that outputs a `tfd.Distribution`.
          batch_size: The batch size to use when checking that the model produces
            the same results as a serialized/deserialized copy.  Default value: 1.