Sub-separable convolution. If separability == 0 it's a separable_conv.

def subseparable_conv(inputs, filters, kernel_size, **kwargs):
  """Sub-separable convolution. If separability == 0 it's a separable_conv."""

  def conv_fn(inputs, filters, kernel_size, **kwargs):
    """Sub-separable convolution, splits into separability-many blocks."""
    separability = None
    if "separability" in kwargs:
      separability = kwargs.pop("separability")
    if separability:
      parts = []
      abs_sep = separability if separability > 0 else -1 * separability
      for split_idx, split in enumerate(tf.split(inputs, abs_sep, axis=3)):
        with tf.variable_scope("part_%d" % split_idx):
          if separability > 0:
            parts.append(
                layers().Conv2D(filters // separability, kernel_size,
                                **kwargs)(split))
          else:
            parts.append(
                layers().SeparableConv2D(filters // abs_sep,
                                         kernel_size, **kwargs)(split))
      if separability > 1:
        result = layers().Conv2D(filters, (1, 1))(tf.concat(parts, axis=3))
      elif abs_sep == 1:  # If we have just one block, return it.
        assert len(parts) == 1
        result = parts[0]
      else:
        result = tf.concat(parts, axis=3)
    else:
      result = layers().SeparableConv2D(filters, kernel_size,
                                        **kwargs)(inputs)
    if separability is not None:
      kwargs["separability"] = separability
    return result

  return conv_internal(conv_fn, inputs, filters, kernel_size, **kwargs)