class_docstr: Masked convolution.
  
  Attributes:
    features: Number of convolution filters.
    kernel_size: Shape of the convolutional kernel. For 1D convolution,
      the kernel size can be passed as an integer. For all other cases, it
      must be a sequence of integers.
    strides: A sequence of `n` integers, representing the inter-window
      strides.
    padding: Either the string `'SAME'`, the string `'VALID'`, or a sequence
      of `n` `(low, high)` integer pairs that give the padding to apply before
      and after each spatial dimension.
    input_dilation: `None`, or a sequence of `n` integers, giving the
      dilation factor to apply in each spatial dimension of `inputs`.
      Convolution with input dilation `d` is equivalent to transposed
      convolution with stride `d`.
    kernel_dilation: `None`, or a sequence of `n` integers, giving the
      dilation factor to apply in each spatial dimension of the convolution
      kernel. Convolution with kernel dilation is also known as 'atrous
      convolution'.
    bias: Whether to add a bias to the output (default: True).
    dtype: The dtype of the computation (default: float32).
    precision: Numerical precision of the computation see `jax.lax.Precision`
      for details.
    kernel_init: Initializer for the convolutional kernel.
    bias_init: Initializer for the bias.
class_docstr_tok: ["Masked","convolution",".","Attributes",":","features",":","Number","of","convolution","filters",".","kernel_size",":","Shape","of","the","convolutional","kernel",".","For","1D","convolution",",","the","kernel","size","can","be","passed","as","an","integer",".","For","all","other","cases",",","it","must","be","a","sequence","of","integers",".","strides",":","A","sequence","of","`","n","`","integers",",","representing","the","inter-window","strides",".","padding",":","Either","the","string","`","'SAME","'","`",",","the","string","`","'VALID","'","`",",","or","a","sequence","of","`","n","`","`","(","low",",","high",")","`","integer","pairs","that","give","the","padding","to","apply","before","and","after","each","spatial","dimension",".","input_dilation",":","`","None","`",",","or","a","sequence","of","`","n","`","integers",",","giving","the","dilation","factor","to","apply","in","each","spatial","dimension","of","`","inputs","`",".","Convolution","with","input","dilation","`","d","`","is","equivalent","to","transposed","convolution","with","stride","`","d","`",".","kernel_dilation",":","`","None","`",",","or","a","sequence","of","`","n","`","integers",",","giving","the","dilation","factor","to","apply","in","each","spatial","dimension","of","the","convolution","kernel",".","Convolution","with","kernel","dilation","is","also","known","as","'atrous","convolution","'",".","bias",":","Whether","to","add","a","bias","to","the","output","(","default",":","True",")",".","dtype",":","The","dtype","of","the","computation","(","default",":","float32",")",".","precision",":","Numerical","precision","of","the","computation","see","`","jax.lax.Precision","`","for","details",".","kernel_init",":","Initializer","for","the","convolutional","kernel",".","bias_init",":","Initializer","for","the","bias","."]
docstr: Applies a *masked* convolution to the inputs.
  
  Args:
    inputs: Input data with dimensions (batch, spatial_dims..., features).
    spatial_shape: Per-input spatial shape of *unpadded* input data with
      dimensions (batch, len(spatial_dims)).
  
  Returns:
    The convolved data and the output spatial size as a tuple.
docstr_tok: ["Applies","a","*","masked","*","convolution","to","the","inputs",".","Args",":","inputs",":","Input","data","with","dimensions","(","batch",",","spatial_dims","...",",","features",")",".","spatial_shape",":","Per-input","spatial","shape","of","*","unpadded","*","input","data","with","dimensions","(","batch",",","len","(","spatial_dims",")",")",".","Returns",":","The","convolved","data","and","the","output","spatial","size","as","a","tuple","."]
code: def __call__(
      self,
      inputs: jnp.ndarray,
      spatial_shape: Optional[jnp.ndarray] = None,
  ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]:
      """Applies a *masked* convolution to the inputs.
  
      Args:
        inputs: Input data with dimensions (batch, spatial_dims..., features).
        spatial_shape: Per-input spatial shape of *unpadded* input data with
          dimensions (batch, len(spatial_dims)).
  
      Returns:
        The convolved data and the output spatial size as a tuple.
      """
      outputs = super(Conv, self).__call__(inputs)
      kernel_size = self.kernel_size
      if isinstance(kernel_size, int):
          kernel_size = (kernel_size,)
  
      if spatial_shape is not None:
          if (
              isinstance(self.padding, str)
              and self.padding.upper() == "SAME"
              and kernel_size != (1,) * len(kernel_size)
          ):
              # In the case of 'SAME' padding the ammounts padded on the left and
              # right depend on the shape of the input (which is dynamic in our case),
              # the stride and the kernel size. In our case this means that each
              # element of the batch might have to be paded in a different way and
              # cannnot be batched.
              # Because the possible number of dynamic pads for a given stride kernel
              # size is finite, it should be possible in theory to run all of them and
              # then select the correct out for masking. The implementation is left
              # for a time when we have a use case for this.
              raise NotImplementedError(
                  "Padding 'SAME' is not supported by masked convolutions."
              )
          spatial_shape = _conv_output_shape(
              spatial_shape=spatial_shape,
              kernel_size=kernel_size,
              input_dilation=self.input_dilation,
              kernel_dilation=self.kernel_dilation,
              strides=self.strides,
              padding=self.padding,
          )
          outputs = apply_spatial_mask(outputs, spatial_shape)
      return outputs, spatial_shape
  
imports: import functools
  from typing import Optional, Tuple, Union, Callable, Any, Sequence, List, Iterable
  import flax.linen
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the Conv obj's underlying member method code with example usage: obj.__call__(inputs, spatial_shape) and returns: outputs,  spatial_shape