class_docstr: U-Net from http://arxiv.org/abs/1505.04597.
  
  Based on:
  https://github.com/NVIDIA/DeepLearningExamples/blob/master/TensorFlow2/Segmentation/UNet_Medical/model/unet.py
  Note that the default configuration `config.padding="VALID"` does only work
  with images that have a certain minimum size (e.g. 128x128 is too small).
  
  Attributes:
    num_classes: Number of classes.
    block_size: Sequence of feature sizes used in UNet blocks.
    padding: Type of padding.
    use_batch_norm: Whether to use batchnorm or not.
class_docstr_tok: ["U-Net","from","http",":","\/\/arxiv.org\/abs\/1505.04597",".","Based","on",":","https",":","\/\/github.com\/NVIDIA\/DeepLearningExamples\/blob\/master\/TensorFlow2\/Segmentation\/UNet_Medical\/model\/unet.py","Note","that","the","default","configuration","`","config.padding=","''","VALID","''","`","does","only","work","with","images","that","have","a","certain","minimum","size","(","e.g",".","128x128","is","too","small",")",".","Attributes",":","num_classes",":","Number","of","classes",".","block_size",":","Sequence","of","feature","sizes","used","in","UNet","blocks",".","padding",":","Type","of","padding",".","use_batch_norm",":","Whether","to","use","batchnorm","or","not","."]
docstr: Applies the UNet model.
docstr_tok: ["Applies","the","UNet","model","."]
code: def __call__(
      self, x: jnp.ndarray, *, train: bool, debug: bool = False
  ) -> jnp.ndarray:
      """Applies the UNet model."""
      del debug
      skip_connections = []
      for i, features in enumerate(self.block_size):
          x, residual = DownsampleBlock(
              features=features,
              padding=self.padding,
              use_batch_norm=self.use_batch_norm,
              name=f"0_down_{i}",
          )(x, train=train)
          skip_connections.append(residual)
      x = BottleneckBlock(
          features=2 * self.block_size[-1],
          padding=self.padding,
          use_batch_norm=self.use_batch_norm,
          name="1_bottleneck",
      )(x, train=train)
  
      *upscaling_features, final_features = self.block_size[::-1]
      for i, features in enumerate(upscaling_features):
          x = UpsampleBlock(
              features=features,
              padding=self.padding,
              use_batch_norm=self.use_batch_norm,
              name=f"2_up_{i}",
          )(x, residual=skip_connections.pop(), train=train)
  
      x = nn_layers.IdentityLayer(name="pre_logits")(x)
      x = OutputBlock(
          features=final_features,
          num_classes=self.num_classes,
          padding=self.padding,
          use_batch_norm=self.use_batch_norm,
          name="output_projection",
      )(x, train=train)
      return x
  
imports: import functools
  from typing import Tuple
  import flax.linen
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.segmentation_model import SegmentationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.model_lib.layers import nn_ops
label_desc: Use image node_id 1 for calling the UNet obj's underlying member method code with example usage: obj.__call__(x) and returns: x