search

tensorflow / lingvo

Lingvo
Apache License 2.0
2.81k stars 445 forks source link

Distributed Shampoo Computing Preconditioners #207

Closed stripathi08 closed 4 years ago

stripathi08 commented 4 years ago

Hi, I'm trying port the distributed shampoo in pytorch but I'm unable to find any code for get/compute ops for preconditioned matrices. Any idea how can these functions be used with other frameworks?

x_ops.compute_preconditioners(
        stats, [-1.0 / (2.0 * r) for r in ranks],
        global_step_int32,
        keys=keys,
        sync=self._synchronous_preconditioning,
        preconditioner_compute_graphdef=self._preconditioner_compute_graphdef
)

Thanks, Shubham

jonathanasdf commented 4 years ago

x_ops links to the C++ custom ops. this specific one is https://github.com/tensorflow/lingvo/blob/74e684489e1e2f214d9df311de009ca2d1355e8b/lingvo/core/ops/preconditioner_op_kernels.cc

stripathi08 commented 4 years ago

Thanks Jonathan, this solves my problem. -Shubham