code: def slicewise(tf_fn,
                xs,
                output_shape=None,
                output_dtype=None,
                splittable_dims=None,
                grad_function=None,
                name=None):
    
    multiple_outputs = isinstance(output_dtype, list)
    output_shapes = output_shape if multiple_outputs else [output_shape]
    output_dtypes = output_dtype if multiple_outputs else [output_dtype]
  
    op = SlicewiseOperation(
        tf_fn,
        xs,
        [convert_to_shape(shape) or xs[0].shape for shape in output_shapes],
        [dtype or xs[0].dtype for dtype in output_dtypes],
        splittable_dims,
        grad_function,
        name=name)
    return tuple(op.outputs) if multiple_outputs else op.outputs[0]
docs: Slice-wise call to any tensorflow function.
  
    The output shape and dtype default to those of the first input.
    splittable_dims is a list of Dimensions which can be split while keeping the
    computation valid.
  
    Args:
      tf_fn: a function taking n tf.Tensors and returning a tf.Tensor
      xs: a list of n Tensors
      output_shape: a Shape (or list of shapes)
      output_dtype: a dtype (or list of dtypes)
      splittable_dims: a list of Dimensions which are ok to split
      grad_function: an optional gradients function.  If None, use tf gradient.
      name: an optional string
  
    Returns:
      a Tensor (or a tuple of Tensors)