docstr: Creates a model that uses a dataset.
  
  Args:
    dataset: A 1D Tensor.
    materialize_dataset: Python `bool`. If True, the dataset is erroneously
      materialized in the initializer.
docstr_tok: ["Creates","a","model","that","uses","a","dataset",".","Args",":","dataset",":","A","1D","Tensor",".","materialize_dataset",":","Python","`","bool","`",".","If","True",",","the","dataset","is","erroneously","materialized","in","the","initializer","."]
code: def __init__(self, dataset, materialize_dataset=False):
      """Creates a model that uses a dataset.
  
      Args:
        dataset: A 1D Tensor.
        materialize_dataset: Python `bool`. If True, the dataset is erroneously
          materialized in the initializer.
      """
      self._distribution = tfd.Sample(
          tfd.Normal(0.0, 1.0), dataset.shape[-1]
      )
      if materialize_dataset:
          tf.convert_to_tensor(dataset)
  
      super(TestModelWithDataset, self).__init__(
          default_event_space_bijector=tfb.Identity(),
          event_shape=self._distribution.event_shape,
          dtype=self._distribution.dtype,
          name="test_model_width_dataset",
          pretty_name="TestModelWithDataset",
          sample_transformations=dict(
              identity=targets.Model.SampleTransformation(
                  fn=lambda x: x,
                  pretty_name="Identity",
                  ground_truth_mean=0.0,
              ),
          ),
      )
  
imports: import atexit
  import functools
  from absl import logging
  import jax
  import numpy
  import tensorflow.compat.v2
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import test_util
  from inference_gym import targets
  from inference_gym.internal import test_util
label_desc: Use image node_id 1 to create a new TestModelWithDataset object from inherited base classes: targets with example: obj = TestModelWithDataset(dataset, materialize_dataset)