Also a part of https://github.com/google/scaaml/issues/131
Also a part of #131