Feature/ Introduce initialization methods for Sinkhorn

[JTT94]

• Examples here: https://colab.research.google.com/drive/1vncmDEr3t6_OKfVC0PJin8PIRBViPyO0?usp=sharing

• Initialization methods stored in /core/initializers.py

  • Each initializer inherits from base class SinkhornInitializer
    • Class is more flexible in case of neural network initialisers in future
    • Currently added:
      • sorting initializer (n=m) for any cost
      • Gaussian initializer for squared ground cost added, only works for point cloud as need access to x
  • Each initializer has methods init_dual_a and init_dual_b

      def init_dual_a(
            self, ot_problem: LinearProblem, lse_mode: bool = True
        ) -> jnp.ndarray:
          """Initialzation for Sinkhorn potential f.
  

  • The base class also holds default behaviour, initialization for 0 and 1 depending on log lse_mode, also handling entries for 0 weights

• Modification to Sinkhorn api,

  • pass instantiated initializer to Sinkhorn

  Sinkhorn(
     lse_mode=lse_mode,
     threshold=threshold,
     norm_error=norm_error,
     inner_iterations=inner_iterations,
     min_iterations=min_iterations,
     max_iterations=max_iterations,
     momentum=momentum_lib.Momentum(start=chg_momentum_from, value=momentum),
     anderson=anderson,
     implicit_diff=implicit_diff,
     parallel_dual_updates=parallel_dual_updates,
     use_danskin=use_danskin,
     potential_initializer=potential_initializer,
     jit=jit
  )
  

  • Can be passed to sinkhorn functional wrapper

  gaus_init = init_lib.GaussianInitializer()
  
  @jax.jit
  def run_sinkhorn_gaus_init(x, y, a=None, b=None):
      sink_kwargs = {'jit': True, 
                  'threshold': 0.001, 
                  'max_iterations': 10**5, 
                  'potential_initializer': gaus_init}
                  
      geom_kwargs = {'epsilon': 0.01}
      geom = PointCloud(x, y, **geom_kwargs)
      out = sinkhorn(geom, a=a, b=b, **sink_kwargs)
      return out
  

  • Initialization strategy will be overridden if init_dual_a/ init_dual_b vectors are passed to solver
  • Addressed this issue Add init to the solver #84 (comment) in tools wrapper

[marcocuturi]

@michalk8 : i had some issues with the webhook to enable recompiling docs on PR. I will look into this soon.

[michalk8]

@JTT94 could you please fix the bibtex file and deal with the rest of the comments? After that, think we can merge.

[JTT94]

Thanks for the feedback, this is good to be merged

[michalk8]

Also just noticed that the initializers are not registered as PyTrees, i.e.

@jax.jit
def test():
    init = ott.core.initializers.SortingInitializer()
    return init

fails. Could you please register them using @jax.tree_util.register_pytree_node_class?

[michalk8]

Thanks a lot @JTT94 , could you please include the last change mentioned here: #98 (comment) ? It's more efficient than doing lax.cond. After that I will happily merge this.

[michalk8]

Thanks @JTT94, LGTM, merging this!

[JTT94]

Thanks @JTT94, LGTM, merging this!

Great! And thanks for your patience with feedback, I appreciate it