def get_linear_lattice_interpolation_fn(lattice_sizes, monotonicities,
                                        output_min, output_max):
  """Returns function which does lattice interpolation.

  Returned function matches lattice_layer.LinearInitializer with corresponding
  parameters.

  Args:
    lattice_sizes: list or tuple of integers which represents lattice sizes.
    monotonicities: monotonicity constraints.
    output_min: minimum output of linear function.
    output_max: maximum output of linear function.

  Returns:
    Function which takes d-dimension point and performs lattice interpolation
    assuming lattice weights are such that lattice represents linear function
    with given output_min and output_max. All monotonic dimesions of this linear
    function cotribute with same weight despite of numer of vertices per
    dimension. All non monotonic dimensions have weight 0.0.
  """

  def linear_interpolation_fn(x):
    """Linear along monotonic dims and 0.0 along non monotonic."""
    result = output_min
    num_monotonic_dims = len(monotonicities) - monotonicities.count(0)
    if num_monotonic_dims == 0:
      local_monotonicities = [1] * len(lattice_sizes)
      num_monotonic_dims = len(lattice_sizes)
    else:
      local_monotonicities = monotonicities

    weight = (output_max - output_min) / num_monotonic_dims
    for i in range(len(x)):
      if local_monotonicities[i]:
        result += x[i] * weight / (lattice_sizes[i] - 1.0)
    return result

  return linear_interpolation_fn